You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@auron.apache.org by ri...@apache.org on 2026/03/19 03:12:58 UTC
(auron) branch dev-v6.0.0-parallel-scan-kdev-build created (now 15b9fcc2)
This is an automated email from the ASF dual-hosted git repository.
richox pushed a change to branch dev-v6.0.0-parallel-scan-kdev-build
in repository https://gitbox.apache.org/repos/asf/auron.git
at 15b9fcc2 add kdev-build-deploy.sh
This branch includes the following new commits:
new 494645c0 Shade more packages (#1121)
new 0eaa812e fix execution error in non-native parquet sink tasks (#1123)
new 83e87e1e fix execute_projected_with_key_rows_output() (#1156)
new b16cc164 prepare version v6.0.0-kwai
new 60328bc9 Fix: keep backward compatibility with celeborn-054This reverts commit 437180b7506bb56da6ad0a77ba6f6e
new d446a839 adapt to celeborn-0.6.0 with some special patchs
new 5f43afda fix exception type
new 10d7d0bb fix empty partition data size
new c4408855 introduce BlazeRuleEngine from spark241kwaiae to spark3
new f7d0b249 supports parallel reading for tables with small files
new 92216fc4 fix possible deadlock in OnHeapSpillManager
new 3fd3ba87 Implement columnar aggregate buffers
new 186718c9 disable string to decimal cast (https://team.corp.kuaishou.com/task/T10924138)
new 08b6fa77 disable max/min + nested params conversion (https://team.corp.kuaishou.com/task/T10892667)
new ea2dd42b supports spark.blaze.enableScanFallback
new 15b9fcc2 add kdev-build-deploy.sh
The 16 revisions listed above as "new" are entirely new to this
repository and will be described in separate emails. The revisions
listed as "add" were already present in the repository and have only
been added to this reference.
(auron) 10/16: supports parallel reading for tables with small files
Posted by ri...@apache.org.
This is an automated email from the ASF dual-hosted git repository.
richox pushed a commit to branch dev-v6.0.0-parallel-scan-kdev-build
in repository https://gitbox.apache.org/repos/asf/auron.git
commit f7d0b24950028da920c3b6ca19ee74f018c753ca
Author: zhangli20 <ri...@gmail.com>
AuthorDate: Wed Dec 24 15:45:33 2025 +0800
supports parallel reading for tables with small files
---
native-engine/blaze-jni-bridge/src/conf.rs | 1 +
native-engine/datafusion-ext-plans/src/lib.rs | 1 +
native-engine/datafusion-ext-plans/src/orc_exec.rs | 76 ++++++------
.../datafusion-ext-plans/src/parquet_exec.rs | 111 +++++++++--------
native-engine/datafusion-ext-plans/src/scan/mod.rs | 1 +
.../datafusion-ext-plans/src/scan/parallel.rs | 138 +++++++++++++++++++++
.../java/org/apache/spark/sql/blaze/BlazeConf.java | 3 +
7 files changed, 242 insertions(+), 89 deletions(-)
diff --git a/native-engine/blaze-jni-bridge/src/conf.rs b/native-engine/blaze-jni-bridge/src/conf.rs
index 99ba7692..c8e77a02 100644
--- a/native-engine/blaze-jni-bridge/src/conf.rs
+++ b/native-engine/blaze-jni-bridge/src/conf.rs
@@ -53,6 +53,7 @@ define_conf!(IntConf, SUGGESTED_BATCH_MEM_SIZE_KWAY_MERGE);
define_conf!(BooleanConf, ORC_FORCE_POSITIONAL_EVOLUTION);
define_conf!(IntConf, UDAF_FALLBACK_NUM_UDAFS_TRIGGER_SORT_AGG);
define_conf!(BooleanConf, PARSE_JSON_ERROR_FALLBACK);
+define_conf!(IntConf, NUM_PARALLEL_SCAN_FILES);
define_conf!(StringConf, NATIVE_LOG_LEVEL);
pub trait BooleanConf {
diff --git a/native-engine/datafusion-ext-plans/src/lib.rs b/native-engine/datafusion-ext-plans/src/lib.rs
index 375338d2..dbde6d08 100644
--- a/native-engine/datafusion-ext-plans/src/lib.rs
+++ b/native-engine/datafusion-ext-plans/src/lib.rs
@@ -19,6 +19,7 @@
#![feature(get_mut_unchecked)]
#![feature(portable_simd)]
#![feature(ptr_as_ref_unchecked)]
+#![feature(unboxed_closures)]
// execution plan implementations
pub mod agg;
diff --git a/native-engine/datafusion-ext-plans/src/orc_exec.rs b/native-engine/datafusion-ext-plans/src/orc_exec.rs
index 074ae31e..3469a49f 100644
--- a/native-engine/datafusion-ext-plans/src/orc_exec.rs
+++ b/native-engine/datafusion-ext-plans/src/orc_exec.rs
@@ -15,11 +15,13 @@
// specific language governing permissions and limitations
// under the License.
-use std::{any::Any, fmt, fmt::Formatter, pin::Pin, sync::Arc};
+use std::{any::Any, fmt, fmt::Formatter, sync::Arc};
use arrow::{datatypes::SchemaRef, error::ArrowError};
use blaze_jni_bridge::{
- conf, conf::BooleanConf, jni_call_static, jni_new_global_ref, jni_new_string,
+ conf,
+ conf::{BooleanConf, IntConf},
+ jni_call_static, jni_new_global_ref, jni_new_string,
};
use bytes::Bytes;
use datafusion::{
@@ -48,7 +50,7 @@ use orc_rust::{
use crate::{
common::execution_context::ExecutionContext,
- scan::{BlazeSchemaMapping, internal_file_reader::InternalFileReader},
+ scan::{BlazeSchemaMapping, internal_file_reader::InternalFileReader, parallel::parallel_scan},
};
/// Execution plan for scanning one or more Orc partitions
@@ -154,33 +156,43 @@ impl ExecutionPlan for OrcExec {
let resource_id = jni_new_string!(&self.fs_resource_id)?;
let fs = jni_call_static!(JniBridge.getResource(resource_id.as_obj()) -> JObject)?;
let fs_provider = Arc::new(FsProvider::new(jni_new_global_ref!(fs.as_obj())?, &io_time));
-
+ let metrics = self.metrics.clone();
let projection = match self.base_config.file_column_projection_indices() {
Some(proj) => proj,
None => (0..self.base_config.file_schema.fields().len()).collect(),
};
let force_positional_evolution = conf::ORC_FORCE_POSITIONAL_EVOLUTION.value()?;
-
- let opener = OrcOpener {
- projection,
- batch_size: batch_size(),
- table_schema: self.base_config.file_schema.clone(),
- fs_provider,
- partition_index: partition,
- metrics: self.metrics.clone(),
- force_positional_evolution,
- };
-
- let file_stream = Box::pin(FileStream::new(
- &self.base_config,
- partition,
- opener,
- exec_ctx.execution_plan_metrics(),
- )?);
-
- let timed_stream = execute_orc_scan(file_stream, exec_ctx.clone())?;
- Ok(exec_ctx.coalesce_with_default_batch_size(timed_stream))
+ let exec_ctx_cloned = exec_ctx.clone();
+ let create_file_scan_config =
+ move |file_scan_config: &FileScanConfig| -> Result<SendableRecordBatchStream> {
+ let opener = OrcOpener {
+ projection: projection.clone(),
+ batch_size: batch_size(),
+ table_schema: file_scan_config.file_schema.clone(),
+ fs_provider: fs_provider.clone(),
+ partition_index: partition,
+ metrics: metrics.clone(),
+ force_positional_evolution,
+ };
+
+ let file_stream = FileStream::new(
+ file_scan_config,
+ partition,
+ opener,
+ exec_ctx_cloned.execution_plan_metrics(),
+ )?;
+ Ok(Box::pin(file_stream))
+ };
+
+ let num_parallel_scan_files = conf::NUM_PARALLEL_SCAN_FILES.value()? as usize;
+ let scan_stream = parallel_scan(
+ exec_ctx.clone(),
+ self.base_config.clone(),
+ create_file_scan_config,
+ num_parallel_scan_files,
+ )?;
+ Ok(exec_ctx.coalesce_with_default_batch_size(scan_stream))
}
fn metrics(&self) -> Option<MetricsSet> {
@@ -192,22 +204,6 @@ impl ExecutionPlan for OrcExec {
}
}
-fn execute_orc_scan(
- mut stream: Pin<Box<FileStream<OrcOpener>>>,
- exec_ctx: Arc<ExecutionContext>,
-) -> Result<SendableRecordBatchStream> {
- Ok(exec_ctx
- .clone()
- .output_with_sender("OrcScan", move |sender| async move {
- sender.exclude_time(exec_ctx.baseline_metrics().elapsed_compute());
- let _timer = exec_ctx.baseline_metrics().elapsed_compute().timer();
- while let Some(batch) = stream.next().await.transpose()? {
- sender.send(batch).await;
- }
- Ok(())
- }))
-}
-
struct OrcOpener {
projection: Vec<usize>,
batch_size: usize,
diff --git a/native-engine/datafusion-ext-plans/src/parquet_exec.rs b/native-engine/datafusion-ext-plans/src/parquet_exec.rs
index a8ab1b35..75679d99 100644
--- a/native-engine/datafusion-ext-plans/src/parquet_exec.rs
+++ b/native-engine/datafusion-ext-plans/src/parquet_exec.rs
@@ -17,12 +17,14 @@
//! Execution plan for reading Parquet files
-use std::{any::Any, fmt, fmt::Formatter, ops::Range, pin::Pin, sync::Arc};
+use std::{any::Any, fmt, fmt::Formatter, ops::Range, sync::Arc};
use arrow::datatypes::SchemaRef;
use arrow_schema::DataType;
use blaze_jni_bridge::{
- conf, conf::BooleanConf, jni_call_static, jni_new_global_ref, jni_new_string,
+ conf,
+ conf::{BooleanConf, IntConf},
+ jni_call_static, jni_new_global_ref, jni_new_string,
};
use bytes::Bytes;
use datafusion::{
@@ -48,7 +50,7 @@ use datafusion::{
};
use datafusion_ext_commons::{batch_size, hadoop_fs::FsProvider};
use fmt::Debug;
-use futures::{FutureExt, StreamExt, future::BoxFuture};
+use futures::{FutureExt, future::BoxFuture};
use itertools::Itertools;
use object_store::ObjectMeta;
use once_cell::sync::OnceCell;
@@ -56,7 +58,10 @@ use parking_lot::Mutex;
use crate::{
common::execution_context::ExecutionContext,
- scan::{BlazeSchemaAdapterFactory, internal_file_reader::InternalFileReader},
+ scan::{
+ BlazeSchemaAdapterFactory, internal_file_reader::InternalFileReader,
+ parallel::parallel_scan,
+ },
};
/// Execution plan for scanning one or more Parquet partitions
@@ -198,46 +203,70 @@ impl ExecutionPlan for ParquetExec {
let _timer = elapsed_compute.timer();
let io_time = exec_ctx.register_timer_metric("io_time");
+ let predicate = self.predicate.clone();
+ let pruning_predicate = self.pruning_predicate.clone();
+ let page_pruning_predicate = self.page_pruning_predicate.clone();
+ let metrics = self.metrics.clone();
+
// get fs object from jni bridge resource
let resource_id = jni_new_string!(&self.fs_resource_id)?;
let fs = jni_call_static!(JniBridge.getResource(resource_id.as_obj()) -> JObject)?;
let fs_provider = Arc::new(FsProvider::new(jni_new_global_ref!(fs.as_obj())?, &io_time));
+ let parquet_file_reader_factory = Arc::new(FsReaderFactory::new(fs_provider.clone()));
let schema_adapter_factory = Arc::new(BlazeSchemaAdapterFactory);
- let projection = match self.base_config.file_column_projection_indices() {
- Some(proj) => proj,
- None => (0..self.base_config.file_schema.fields().len()).collect(),
+ let projection: Arc<[usize]> = match self.base_config.file_column_projection_indices() {
+ Some(proj) => proj.into(),
+ None => (0..self.base_config.file_schema.fields().len())
+ .collect::<Vec<_>>()
+ .into(),
};
let page_filtering_enabled = conf::PARQUET_ENABLE_PAGE_FILTERING.value()?;
let bloom_filter_enabled = conf::PARQUET_ENABLE_BLOOM_FILTER.value()?;
+ let base_config = self.base_config.clone();
+
+ let exec_ctx_cloned = exec_ctx.clone();
+ let create_file_stream_fn =
+ move |file_scan_config: &FileScanConfig| -> Result<SendableRecordBatchStream> {
+ let opener = ParquetOpener {
+ partition_index: partition,
+ projection: projection.clone(),
+ batch_size: batch_size(),
+ limit: file_scan_config.limit,
+ predicate: predicate.clone(),
+ pruning_predicate: pruning_predicate.clone(),
+ page_pruning_predicate: page_pruning_predicate.clone(),
+ table_schema: file_scan_config.file_schema.clone(),
+ metadata_size_hint: None,
+ metrics: metrics.clone(),
+ parquet_file_reader_factory: parquet_file_reader_factory.clone(),
+ pushdown_filters: page_filtering_enabled,
+ reorder_filters: page_filtering_enabled,
+ enable_page_index: page_filtering_enabled,
+ enable_bloom_filter: bloom_filter_enabled,
+ schema_adapter_factory: schema_adapter_factory.clone(),
+ };
+ let mut file_stream = FileStream::new(
+ file_scan_config,
+ exec_ctx_cloned.partition_id(),
+ opener,
+ exec_ctx_cloned.execution_plan_metrics(),
+ )?;
+ if conf::IGNORE_CORRUPTED_FILES.value()? {
+ file_stream = file_stream.with_on_error(OnError::Skip);
+ }
+ Ok(Box::pin(file_stream))
+ };
- let opener = ParquetOpener {
- partition_index: partition,
- projection: Arc::from(projection),
- batch_size: batch_size(),
- limit: self.base_config.limit,
- predicate: self.predicate.clone(),
- pruning_predicate: self.pruning_predicate.clone(),
- page_pruning_predicate: self.page_pruning_predicate.clone(),
- table_schema: self.base_config.file_schema.clone(),
- metadata_size_hint: None,
- metrics: self.metrics.clone(),
- parquet_file_reader_factory: Arc::new(FsReaderFactory::new(fs_provider)),
- pushdown_filters: page_filtering_enabled,
- reorder_filters: page_filtering_enabled,
- enable_page_index: page_filtering_enabled,
- enable_bloom_filter: bloom_filter_enabled,
- schema_adapter_factory,
- };
-
- let mut file_stream = FileStream::new(&self.base_config, partition, opener, &self.metrics)?;
- if conf::IGNORE_CORRUPTED_FILES.value()? {
- file_stream = file_stream.with_on_error(OnError::Skip);
- }
-
- let timed_stream = execute_parquet_scan(Box::pin(file_stream), exec_ctx.clone())?;
- Ok(exec_ctx.coalesce_with_default_batch_size(timed_stream))
+ let num_parallel_scan_files = conf::NUM_PARALLEL_SCAN_FILES.value()? as usize;
+ let scan_stream = parallel_scan(
+ exec_ctx.clone(),
+ base_config,
+ create_file_stream_fn,
+ num_parallel_scan_files,
+ )?;
+ Ok(exec_ctx.coalesce_with_default_batch_size(scan_stream))
}
fn metrics(&self) -> Option<MetricsSet> {
@@ -249,22 +278,6 @@ impl ExecutionPlan for ParquetExec {
}
}
-fn execute_parquet_scan(
- mut stream: Pin<Box<FileStream<ParquetOpener>>>,
- exec_ctx: Arc<ExecutionContext>,
-) -> Result<SendableRecordBatchStream> {
- Ok(exec_ctx
- .clone()
- .output_with_sender("ParquetScan", move |sender| async move {
- sender.exclude_time(exec_ctx.baseline_metrics().elapsed_compute());
- let _timer = exec_ctx.baseline_metrics().elapsed_compute().timer();
- while let Some(batch) = stream.next().await.transpose()? {
- sender.send(batch).await;
- }
- Ok(())
- }))
-}
-
#[derive(Clone)]
pub struct FsReaderFactory {
fs_provider: Arc<FsProvider>,
diff --git a/native-engine/datafusion-ext-plans/src/scan/mod.rs b/native-engine/datafusion-ext-plans/src/scan/mod.rs
index ba89500b..ad4b69cd 100644
--- a/native-engine/datafusion-ext-plans/src/scan/mod.rs
+++ b/native-engine/datafusion-ext-plans/src/scan/mod.rs
@@ -25,6 +25,7 @@ use datafusion::{
use datafusion_ext_commons::df_execution_err;
pub mod internal_file_reader;
+pub mod parallel;
#[derive(Debug)]
pub struct BlazeSchemaAdapterFactory;
diff --git a/native-engine/datafusion-ext-plans/src/scan/parallel.rs b/native-engine/datafusion-ext-plans/src/scan/parallel.rs
new file mode 100644
index 00000000..fc6f38cf
--- /dev/null
+++ b/native-engine/datafusion-ext-plans/src/scan/parallel.rs
@@ -0,0 +1,138 @@
+// Copyright 2022 The Blaze Authors
+//
+// Licensed 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 std::sync::Arc;
+
+use blaze_jni_bridge::is_task_running;
+use datafusion::{
+ common::DataFusionError, datasource::physical_plan::FileScanConfig, error::Result,
+ execution::SendableRecordBatchStream,
+};
+use datafusion_ext_commons::df_execution_err;
+use futures_util::StreamExt;
+
+use crate::common::execution_context::ExecutionContext;
+
+pub fn parallel_scan(
+ exec_ctx: Arc<ExecutionContext>,
+ file_scan_config: FileScanConfig,
+ create_file_stream_fn: impl Fn(&FileScanConfig) -> Result<SendableRecordBatchStream>
+ + Send
+ + 'static,
+ num_parallel_files: usize,
+) -> Result<SendableRecordBatchStream> {
+ let partition_files = file_scan_config
+ .file_groups
+ .iter()
+ .flatten()
+ .cloned()
+ .collect::<Vec<_>>();
+
+ if num_parallel_files <= 0 {
+ return df_execution_err!("num_parallel_files must be positive, got {num_parallel_files}");
+ }
+
+ // no parallel
+ if num_parallel_files == 1 {
+ let mut file_stream = create_file_stream_fn(&file_scan_config)?;
+ let stream = exec_ctx
+ .clone()
+ .output_with_sender("FileScan", move |sender| async move {
+ let elapsed_compute = exec_ctx.baseline_metrics().elapsed_compute().clone();
+ let _timer = elapsed_compute.timer();
+ sender.exclude_time(&elapsed_compute);
+
+ while let Some(batch) = file_stream.next().await.transpose()? {
+ sender.send(batch).await;
+ }
+ Ok(())
+ });
+ return Ok(stream);
+ }
+
+ let stream = exec_ctx
+ .clone()
+ .output_with_sender("FileScan", move |sender| async move {
+ let elapsed_compute = exec_ctx.baseline_metrics().elapsed_compute().clone();
+ let _timer = elapsed_compute.timer();
+ sender.exclude_time(&elapsed_compute);
+
+ let (file_streams_tx, mut file_streams_rx) =
+ tokio::sync::mpsc::channel(num_parallel_files - 1);
+ let exec_ctx_cloned = exec_ctx.clone();
+
+ // create file streams for each file
+ let mut file_streams: Vec<SendableRecordBatchStream> = partition_files
+ .into_iter()
+ .map(|partition_file| {
+ let mut file_scan_config = file_scan_config.clone();
+ let partition_id = exec_ctx_cloned.partition_id();
+ file_scan_config.file_groups = vec![vec![]; file_scan_config.file_groups.len()];
+ file_scan_config.file_groups[partition_id] = vec![partition_file];
+ create_file_stream_fn(&file_scan_config)
+ })
+ .collect::<Result<_>>()?;
+
+ // read first batch without parallelism to avoid latency of later operators
+ if !file_streams.is_empty() {
+ if let Some(batch) = file_streams[0].next().await.transpose()? {
+ sender.send(batch).await;
+ }
+ }
+
+ // read rest file streams in parallel
+ let handle = tokio::spawn(async move {
+ for mut file_stream in file_streams {
+ let (tx, mut rx) = tokio::sync::mpsc::channel(1);
+ let handle = tokio::spawn(async move {
+ while is_task_running()
+ && let Some(batch) = file_stream.next().await.transpose()?
+ {
+ tx.send(batch).await.or_else(|e| df_execution_err!("{e}"))?;
+ }
+ Ok::<_, DataFusionError>(())
+ });
+ let eager_stream = exec_ctx_cloned.output_with_sender(
+ "FileScan.File",
+ move |sender| async move {
+ while is_task_running()
+ && let Some(batch) = rx.recv().await
+ {
+ sender.send(batch).await;
+ }
+ handle.await.or_else(|e| df_execution_err!("{e}"))??;
+ Ok(())
+ },
+ );
+
+ file_streams_tx
+ .send(eager_stream)
+ .await
+ .or_else(|e| df_execution_err!("{e}"))?;
+ }
+ Ok::<_, DataFusionError>(())
+ });
+
+ while let Some(mut stream) = file_streams_rx.recv().await {
+ while let Some(batch) = stream.next().await.transpose()? {
+ sender.send(batch).await;
+ }
+ }
+ handle
+ .await
+ .or_else(|e| df_execution_err!("failed to join task handle: {e}"))??;
+ Ok(())
+ });
+ Ok(stream)
+}
diff --git a/spark-extension/src/main/java/org/apache/spark/sql/blaze/BlazeConf.java b/spark-extension/src/main/java/org/apache/spark/sql/blaze/BlazeConf.java
index e655f79e..5136b8b5 100644
--- a/spark-extension/src/main/java/org/apache/spark/sql/blaze/BlazeConf.java
+++ b/spark-extension/src/main/java/org/apache/spark/sql/blaze/BlazeConf.java
@@ -112,6 +112,9 @@ public enum BlazeConf {
ORC_FORCE_POSITIONAL_EVOLUTION("spark.blaze.orc.force.positional.evolution", false),
+ // number of parallel scan files
+ NUM_PARALLEL_SCAN_FILES("spark.blaze.numParallelScanFiles", 4),
+
NATIVE_LOG_LEVEL("spark.blaze.native.log.level", "info");
public final String key;
(auron) 08/16: fix empty partition data size
Posted by ri...@apache.org.
This is an automated email from the ASF dual-hosted git repository.
richox pushed a commit to branch dev-v6.0.0-parallel-scan-kdev-build
in repository https://gitbox.apache.org/repos/asf/auron.git
commit 10d7d0bbe144f05a145ea63d8a2f3463038849c1
Author: zhuangxian <zh...@kuaishou.com>
AuthorDate: Wed Dec 10 13:26:34 2025 +0000
fix empty partition data size
KDev_MR_link：https://ksurl.cn/yez9gIJf
---
pom.xml | 2 +-
.../celeborn/BlazeCelebornShuffleWriter.scala | 20 ++++++++++++++++++--
2 files changed, 19 insertions(+), 3 deletions(-)
diff --git a/pom.xml b/pom.xml
index 57140bf3..0301c25f 100644
--- a/pom.xml
+++ b/pom.xml
@@ -416,7 +416,7 @@
<id>spark-3.5</id>
<properties>
<shimName>spark-3.5</shimName>
- <pkgSuffix>-kwai-adapt-cele060</pkgSuffix>
+ <pkgSuffix>-kwai-adapt-cele060-fix-non-partition-len</pkgSuffix>
<shimPkg>spark-extension-shims-spark3</shimPkg>
<javaVersion>1.8</javaVersion>
<scalaVersion>2.12</scalaVersion>
diff --git a/spark-extension-shims-spark3/src/main/scala/org/apache/spark/sql/execution/blaze/shuffle/celeborn/BlazeCelebornShuffleWriter.scala b/spark-extension-shims-spark3/src/main/scala/org/apache/spark/sql/execution/blaze/shuffle/celeborn/BlazeCelebornShuffleWriter.scala
index 2c636307..c5c4df76 100644
--- a/spark-extension-shims-spark3/src/main/scala/org/apache/spark/sql/execution/blaze/shuffle/celeborn/BlazeCelebornShuffleWriter.scala
+++ b/spark-extension-shims-spark3/src/main/scala/org/apache/spark/sql/execution/blaze/shuffle/celeborn/BlazeCelebornShuffleWriter.scala
@@ -16,7 +16,7 @@
package org.apache.spark.sql.execution.blaze.shuffle.celeborn
import org.apache.celeborn.client.ShuffleClient
-import org.apache.spark.TaskContext
+import org.apache.spark.{TaskContext, SparkEnv}
import org.apache.spark.scheduler.MapStatus
import org.apache.spark.shuffle.ShuffleHandle
import org.apache.spark.shuffle.ShuffleWriteMetricsReporter
@@ -24,6 +24,7 @@ import org.apache.spark.shuffle.ShuffleWriter
import org.apache.spark.shuffle.celeborn.CelebornShuffleHandle
import org.apache.spark.shuffle.celeborn.ExecutorShuffleIdTracker
import org.apache.spark.shuffle.celeborn.SparkUtils
+import org.apache.spark.sql.blaze.Shims
import org.apache.spark.sql.execution.blaze.shuffle.BlazeRssShuffleWriterBase
import org.apache.spark.sql.execution.blaze.shuffle.RssPartitionWriterBase
import org.blaze.sparkver
@@ -68,4 +69,19 @@ class BlazeCelebornShuffleWriter[K, V](
celebornShuffleWriter.write(Iterator.empty) // force flush
celebornShuffleWriter.stop(success)
}
-}
+
+ // Override stop to use partition length map directly instead of rssStop's mapStatus
+ // because celeborn writer doesn't populate partition sizes correctly when using native writer
+ override def stop(success: Boolean): Option[MapStatus] = {
+ if (!success) {
+ return None
+ }
+
+ // Always use getPartitionLengthMap for Celeborn to get correct partition sizes
+ val blockManagerId = SparkEnv.get.blockManager.shuffleServerId
+ Some(Shims.get.getMapStatus(
+ blockManagerId,
+ celebornPartitionWriter.getPartitionLengthMap,
+ taskContext.partitionId()))
+ }
+}
\ No newline at end of file
(auron) 04/16: prepare version v6.0.0-kwai
Posted by ri...@apache.org.
This is an automated email from the ASF dual-hosted git repository.
richox pushed a commit to branch dev-v6.0.0-parallel-scan-kdev-build
in repository https://gitbox.apache.org/repos/asf/auron.git
commit b16cc1647c2fb0bb7f172955966496572f3ae3b2
Author: zhangli20 <zh...@kuaishou.com>
AuthorDate: Thu Apr 10 14:27:04 2025 +0800
prepare version v6.0.0-kwai
---
.gitignore | 3 -
Cargo.lock | 1020 +++++++++++---------
Cargo.toml | 40 +-
build-native.sh | 2 +-
dev/mvn-build-helper/assembly/pom.xml | 2 +-
pom.xml | 29 +-
spark-extension-shims-spark241kwaiae/pom.xml | 66 ++
.../org/apache/spark/sql/blaze/ShimsImpl.scala | 813 ++++++++++++++++
.../blaze/kwai/BlazeOperatorMetricsCollector.scala | 130 +++
.../kwai/KwaiPrivilegedHDFSBlockManager.scala | 51 +
.../execution/blaze/plan/ConvertToNativeExec.scala | 12 +-
.../sql/execution/blaze/plan/NativeAggExec.scala | 44 +
.../blaze/plan/NativeBroadcastExchangeExec.scala | 53 +
.../execution/blaze/plan/NativeExpandExec.scala | 19 +-
.../execution/blaze/plan/NativeFilterExec.scala | 15 +-
.../execution/blaze/plan/NativeGenerateExec.scala | 32 +
.../blaze/plan/NativeGlobalLimitExec.scala | 14 +-
.../blaze/plan/NativeLocalLimitExec.scala | 14 +-
.../execution/blaze/plan/NativeOrcScanExec.scala | 16 +-
.../NativeParquetInsertIntoHiveTableExec.scala | 151 +++
.../blaze/plan/NativeParquetScanExec.scala | 16 +-
.../blaze/plan/NativeParquetSinkExec.scala | 57 ++
.../blaze/plan/NativePartialTakeOrderedExec.scala | 26 +-
.../execution/blaze/plan/NativeProjectExec.scala | 15 +-
.../blaze/plan/NativeRenameColumnsExec.scala | 14 +-
.../blaze/plan/NativeShuffleExchangeExec.scala | 236 +++++
.../sql/execution/blaze/plan/NativeSortExec.scala | 17 +-
.../blaze/plan/NativeSortMergeJoinExec.scala | 39 +
.../blaze/plan/NativeTakeOrderedExec.scala | 18 +-
.../sql/execution/blaze/plan/NativeUnionExec.scala | 17 +-
.../execution/blaze/plan/NativeWindowExec.scala | 32 +
.../shuffle/BlazeBlockStoreShuffleReader.scala | 87 ++
.../shuffle/BlazeRssBlockStoreShuffleReader.scala | 154 +++
.../blaze/shuffle/BlazeRssShuffleManager.scala | 130 +++
.../blaze/shuffle/BlazeRssShuffleWriter.scala | 49 +
.../blaze/shuffle/BlazeShuffleManager.scala | 127 +++
.../blaze/shuffle/BlazeShuffleWriter.scala | 10 +-
.../blaze/shuffle/RssPartitionWriter.scala | 86 ++
.../joins/blaze/plan/NativeBroadcastJoinExec.scala | 59 ++
.../blaze/plan/NativeShuffledHashJoinExec.scala | 54 ++
.../sql/execution/ui/BlazeSQLPlanCollector.scala | 301 ++++++
.../shuffle/uniffle/RssShuffleHandleWrapper.scala | 6 +-
.../org/apache/spark/sql/blaze/ShimsImpl.scala | 68 +-
.../kwai/KwaiPrivilegedHDFSBlockManager.scala | 51 +
.../blaze/shuffle/BlazeRssShuffleManagerBase.scala | 1 +
.../blaze/shuffle/BlazeShuffleManager.scala | 11 +-
.../celeborn/BlazeCelebornShuffleManager.scala | 3 +-
.../uniffle/BlazeUniffleShuffleManager.scala | 3 +-
.../apache/spark/sql/blaze/BaseBlazeSQLSuite.scala | 5 +-
.../spark/sql/blaze/BlazeFunctionSuite.scala | 29 -
.../apache/spark/sql/blaze/BlazeConverters.scala | 19 +-
.../sql/blaze/BlazeSparkSessionExtension.scala | 2 +-
.../apache/spark/sql/blaze/NativeConverters.scala | 27 +-
.../scala/org/apache/spark/sql/blaze/Shims.scala | 4 +-
.../blaze/plan/NativeBroadcastExchangeBase.scala | 6 +
.../blaze/plan/NativeShuffledHashJoinBase.scala | 2 +-
.../blaze/plan/NativePaimonTableScanExec.scala | 8 +
57 files changed, 3669 insertions(+), 646 deletions(-)
diff --git a/.gitignore b/.gitignore
index f4cb2013..f5f0d128 100644
--- a/.gitignore
+++ b/.gitignore
@@ -48,9 +48,6 @@ target/
# Protoc generated src
spark-extension/src/generated-sources
-# exclude spark241kwaiae shim module
-spark-extension-shims-spark241kwaiae
-
# exclude native engine build lib directory
native-engine/_build
diff --git a/Cargo.lock b/Cargo.lock
index ca338307..9bed850f 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -13,22 +13,22 @@ dependencies = [
[[package]]
name = "adler2"
-version = "2.0.0"
+version = "2.0.1"
source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "512761e0bb2578dd7380c6baaa0f4ce03e84f95e960231d1dec8bf4d7d6e2627"
+checksum = "320119579fcad9c21884f5c4861d16174d0e06250625266f50fe6898340abefa"
[[package]]
name = "ahash"
-version = "0.8.11"
+version = "0.8.12"
source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "e89da841a80418a9b391ebaea17f5c112ffaaa96f621d2c285b5174da76b9011"
+checksum = "5a15f179cd60c4584b8a8c596927aadc462e27f2ca70c04e0071964a73ba7a75"
dependencies = [
"cfg-if",
"const-random",
- "getrandom 0.2.15",
+ "getrandom 0.3.3",
"once_cell",
"version_check",
- "zerocopy 0.7.35",
+ "zerocopy",
]
[[package]]
@@ -42,9 +42,9 @@ dependencies = [
[[package]]
name = "aligned-vec"
-version = "0.6.1"
+version = "0.6.4"
source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "7e0966165eaf052580bd70eb1b32cb3d6245774c0104d1b2793e9650bf83b52a"
+checksum = "dc890384c8602f339876ded803c97ad529f3842aba97f6392b3dba0dd171769b"
dependencies = [
"equator",
]
@@ -87,9 +87,9 @@ dependencies = [
[[package]]
name = "anyhow"
-version = "1.0.95"
+version = "1.0.98"
source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "34ac096ce696dc2fcabef30516bb13c0a68a11d30131d3df6f04711467681b04"
+checksum = "e16d2d3311acee920a9eb8d33b8cbc1787ce4a264e85f964c2404b969bdcd487"
[[package]]
name = "arrayref"
@@ -106,7 +106,7 @@ checksum = "7c02d123df017efcdfbd739ef81735b36c5ba83ec3c59c80a9d7ecc718f92e50"
[[package]]
name = "arrow"
version = "53.0.0"
-source = "git+https://github.com/blaze-init/arrow-rs.git?rev=1fbdc9f980aa34c7758da9db79a64248f24daf1e#1fbdc9f980aa34c7758da9db79a64248f24daf1e"
+source = "git+https://github.com/blaze-init/arrow-rs.git?rev=94ef6f608254484c0ecfc33fcef19891c4f87612#94ef6f608254484c0ecfc33fcef19891c4f87612"
dependencies = [
"arrow-arith",
"arrow-array",
@@ -126,7 +126,7 @@ dependencies = [
[[package]]
name = "arrow-arith"
version = "53.0.0"
-source = "git+https://github.com/blaze-init/arrow-rs.git?rev=1fbdc9f980aa34c7758da9db79a64248f24daf1e#1fbdc9f980aa34c7758da9db79a64248f24daf1e"
+source = "git+https://github.com/blaze-init/arrow-rs.git?rev=94ef6f608254484c0ecfc33fcef19891c4f87612#94ef6f608254484c0ecfc33fcef19891c4f87612"
dependencies = [
"arrow-array",
"arrow-buffer",
@@ -140,7 +140,7 @@ dependencies = [
[[package]]
name = "arrow-array"
version = "53.0.0"
-source = "git+https://github.com/blaze-init/arrow-rs.git?rev=1fbdc9f980aa34c7758da9db79a64248f24daf1e#1fbdc9f980aa34c7758da9db79a64248f24daf1e"
+source = "git+https://github.com/blaze-init/arrow-rs.git?rev=94ef6f608254484c0ecfc33fcef19891c4f87612#94ef6f608254484c0ecfc33fcef19891c4f87612"
dependencies = [
"ahash",
"arrow-buffer",
@@ -156,7 +156,7 @@ dependencies = [
[[package]]
name = "arrow-buffer"
version = "53.0.0"
-source = "git+https://github.com/blaze-init/arrow-rs.git?rev=1fbdc9f980aa34c7758da9db79a64248f24daf1e#1fbdc9f980aa34c7758da9db79a64248f24daf1e"
+source = "git+https://github.com/blaze-init/arrow-rs.git?rev=94ef6f608254484c0ecfc33fcef19891c4f87612#94ef6f608254484c0ecfc33fcef19891c4f87612"
dependencies = [
"bytes 1.10.1",
"half",
@@ -166,7 +166,7 @@ dependencies = [
[[package]]
name = "arrow-cast"
version = "53.0.0"
-source = "git+https://github.com/blaze-init/arrow-rs.git?rev=1fbdc9f980aa34c7758da9db79a64248f24daf1e#1fbdc9f980aa34c7758da9db79a64248f24daf1e"
+source = "git+https://github.com/blaze-init/arrow-rs.git?rev=94ef6f608254484c0ecfc33fcef19891c4f87612#94ef6f608254484c0ecfc33fcef19891c4f87612"
dependencies = [
"arrow-array",
"arrow-buffer",
@@ -186,7 +186,7 @@ dependencies = [
[[package]]
name = "arrow-csv"
version = "53.0.0"
-source = "git+https://github.com/blaze-init/arrow-rs.git?rev=1fbdc9f980aa34c7758da9db79a64248f24daf1e#1fbdc9f980aa34c7758da9db79a64248f24daf1e"
+source = "git+https://github.com/blaze-init/arrow-rs.git?rev=94ef6f608254484c0ecfc33fcef19891c4f87612#94ef6f608254484c0ecfc33fcef19891c4f87612"
dependencies = [
"arrow-array",
"arrow-buffer",
@@ -204,7 +204,7 @@ dependencies = [
[[package]]
name = "arrow-data"
version = "53.0.0"
-source = "git+https://github.com/blaze-init/arrow-rs.git?rev=1fbdc9f980aa34c7758da9db79a64248f24daf1e#1fbdc9f980aa34c7758da9db79a64248f24daf1e"
+source = "git+https://github.com/blaze-init/arrow-rs.git?rev=94ef6f608254484c0ecfc33fcef19891c4f87612#94ef6f608254484c0ecfc33fcef19891c4f87612"
dependencies = [
"arrow-buffer",
"arrow-schema",
@@ -215,7 +215,7 @@ dependencies = [
[[package]]
name = "arrow-ipc"
version = "53.0.0"
-source = "git+https://github.com/blaze-init/arrow-rs.git?rev=1fbdc9f980aa34c7758da9db79a64248f24daf1e#1fbdc9f980aa34c7758da9db79a64248f24daf1e"
+source = "git+https://github.com/blaze-init/arrow-rs.git?rev=94ef6f608254484c0ecfc33fcef19891c4f87612#94ef6f608254484c0ecfc33fcef19891c4f87612"
dependencies = [
"arrow-array",
"arrow-buffer",
@@ -229,7 +229,7 @@ dependencies = [
[[package]]
name = "arrow-json"
version = "53.0.0"
-source = "git+https://github.com/blaze-init/arrow-rs.git?rev=1fbdc9f980aa34c7758da9db79a64248f24daf1e#1fbdc9f980aa34c7758da9db79a64248f24daf1e"
+source = "git+https://github.com/blaze-init/arrow-rs.git?rev=94ef6f608254484c0ecfc33fcef19891c4f87612#94ef6f608254484c0ecfc33fcef19891c4f87612"
dependencies = [
"arrow-array",
"arrow-buffer",
@@ -248,7 +248,7 @@ dependencies = [
[[package]]
name = "arrow-ord"
version = "53.0.0"
-source = "git+https://github.com/blaze-init/arrow-rs.git?rev=1fbdc9f980aa34c7758da9db79a64248f24daf1e#1fbdc9f980aa34c7758da9db79a64248f24daf1e"
+source = "git+https://github.com/blaze-init/arrow-rs.git?rev=94ef6f608254484c0ecfc33fcef19891c4f87612#94ef6f608254484c0ecfc33fcef19891c4f87612"
dependencies = [
"arrow-array",
"arrow-buffer",
@@ -262,7 +262,7 @@ dependencies = [
[[package]]
name = "arrow-row"
version = "53.0.0"
-source = "git+https://github.com/blaze-init/arrow-rs.git?rev=1fbdc9f980aa34c7758da9db79a64248f24daf1e#1fbdc9f980aa34c7758da9db79a64248f24daf1e"
+source = "git+https://github.com/blaze-init/arrow-rs.git?rev=94ef6f608254484c0ecfc33fcef19891c4f87612#94ef6f608254484c0ecfc33fcef19891c4f87612"
dependencies = [
"ahash",
"arrow-array",
@@ -275,16 +275,16 @@ dependencies = [
[[package]]
name = "arrow-schema"
version = "53.0.0"
-source = "git+https://github.com/blaze-init/arrow-rs.git?rev=1fbdc9f980aa34c7758da9db79a64248f24daf1e#1fbdc9f980aa34c7758da9db79a64248f24daf1e"
+source = "git+https://github.com/blaze-init/arrow-rs.git?rev=94ef6f608254484c0ecfc33fcef19891c4f87612#94ef6f608254484c0ecfc33fcef19891c4f87612"
dependencies = [
- "bitflags 2.6.0",
+ "bitflags 2.9.1",
"serde",
]
[[package]]
name = "arrow-select"
version = "53.0.0"
-source = "git+https://github.com/blaze-init/arrow-rs.git?rev=1fbdc9f980aa34c7758da9db79a64248f24daf1e#1fbdc9f980aa34c7758da9db79a64248f24daf1e"
+source = "git+https://github.com/blaze-init/arrow-rs.git?rev=94ef6f608254484c0ecfc33fcef19891c4f87612#94ef6f608254484c0ecfc33fcef19891c4f87612"
dependencies = [
"ahash",
"arrow-array",
@@ -297,7 +297,7 @@ dependencies = [
[[package]]
name = "arrow-string"
version = "53.0.0"
-source = "git+https://github.com/blaze-init/arrow-rs.git?rev=1fbdc9f980aa34c7758da9db79a64248f24daf1e#1fbdc9f980aa34c7758da9db79a64248f24daf1e"
+source = "git+https://github.com/blaze-init/arrow-rs.git?rev=94ef6f608254484c0ecfc33fcef19891c4f87612#94ef6f608254484c0ecfc33fcef19891c4f87612"
dependencies = [
"arrow-array",
"arrow-buffer",
@@ -312,11 +312,11 @@ dependencies = [
[[package]]
name = "async-compression"
-version = "0.4.18"
+version = "0.4.19"
source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "df895a515f70646414f4b45c0b79082783b80552b373a68283012928df56f522"
+checksum = "06575e6a9673580f52661c92107baabffbf41e2141373441cbcdc47cb733003c"
dependencies = [
- "bzip2",
+ "bzip2 0.5.2",
"flate2",
"futures-core",
"futures-io",
@@ -336,7 +336,7 @@ checksum = "e539d3fca749fcee5236ab05e93a52867dd549cc157c8cb7f99595f3cedffdb5"
dependencies = [
"proc-macro2",
"quote",
- "syn 2.0.91",
+ "syn 2.0.104",
]
[[package]]
@@ -356,15 +356,15 @@ checksum = "1505bd5d3d116872e7271a6d4e16d81d0c8570876c8de68093a09ac269d8aac0"
[[package]]
name = "autocfg"
-version = "1.4.0"
+version = "1.5.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "ace50bade8e6234aa140d9a2f552bbee1db4d353f69b8217bc503490fc1a9f26"
+checksum = "c08606f8c3cbf4ce6ec8e28fb0014a2c086708fe954eaa885384a6165172e7e8"
[[package]]
name = "aws-lc-rs"
-version = "1.13.1"
+version = "1.13.3"
source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "93fcc8f365936c834db5514fc45aee5b1202d677e6b40e48468aaaa8183ca8c7"
+checksum = "5c953fe1ba023e6b7730c0d4b031d06f267f23a46167dcbd40316644b10a17ba"
dependencies = [
"aws-lc-sys",
"zeroize",
@@ -372,9 +372,9 @@ dependencies = [
[[package]]
name = "aws-lc-sys"
-version = "0.29.0"
+version = "0.30.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "61b1d86e7705efe1be1b569bab41d4fa1e14e220b60a160f78de2db687add079"
+checksum = "dbfd150b5dbdb988bcc8fb1fe787eb6b7ee6180ca24da683b61ea5405f3d43ff"
dependencies = [
"bindgen",
"cc",
@@ -385,9 +385,9 @@ dependencies = [
[[package]]
name = "backtrace"
-version = "0.3.74"
+version = "0.3.75"
source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "8d82cb332cdfaed17ae235a638438ac4d4839913cc2af585c3c6746e8f8bee1a"
+checksum = "6806a6321ec58106fea15becdad98371e28d92ccbc7c8f1b3b6dd724fe8f1002"
dependencies = [
"addr2line",
"cfg-if",
@@ -395,7 +395,7 @@ dependencies = [
"miniz_oxide",
"object",
"rustc-demangle",
- "windows-targets",
+ "windows-targets 0.52.6",
]
[[package]]
@@ -423,7 +423,7 @@ version = "0.69.5"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "271383c67ccabffb7381723dea0672a673f292304fcb45c01cc648c7a8d58088"
dependencies = [
- "bitflags 2.6.0",
+ "bitflags 2.9.1",
"cexpr",
"clang-sys",
"itertools 0.12.1",
@@ -436,7 +436,7 @@ dependencies = [
"regex",
"rustc-hash",
"shlex",
- "syn 2.0.91",
+ "syn 2.0.104",
"which",
]
@@ -448,9 +448,9 @@ checksum = "bef38d45163c2f1dde094a7dfd33ccf595c92905c8f8f4fdc18d06fb1037718a"
[[package]]
name = "bitflags"
-version = "2.6.0"
+version = "2.9.1"
source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "b048fb63fd8b5923fc5aa7b340d8e156aec7ec02f0c78fa8a6ddc2613f6f71de"
+checksum = "1b8e56985ec62d17e9c1001dc89c88ecd7dc08e47eba5ec7c29c7b5eeecde967"
[[package]]
name = "bitvec"
@@ -475,9 +475,9 @@ dependencies = [
[[package]]
name = "blake3"
-version = "1.5.5"
+version = "1.8.2"
source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "b8ee0c1824c4dea5b5f81736aff91bae041d2c07ee1192bec91054e10e3e601e"
+checksum = "3888aaa89e4b2a40fca9848e400f6a658a5a3978de7be858e209cafa8be9a4a0"
dependencies = [
"arrayref",
"arrayvec",
@@ -565,9 +565,9 @@ dependencies = [
[[package]]
name = "brotli-decompressor"
-version = "4.0.1"
+version = "4.0.3"
source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "9a45bd2e4095a8b518033b128020dd4a55aab1c0a381ba4404a472630f4bc362"
+checksum = "a334ef7c9e23abf0ce748e8cd309037da93e606ad52eb372e4ce327a0dcfbdfd"
dependencies = [
"alloc-no-stdlib",
"alloc-stdlib",
@@ -581,9 +581,9 @@ checksum = "46c5e41b57b8bba42a04676d81cb89e9ee8e859a1a66f80a5a72e1cb76b34d43"
[[package]]
name = "bytemuck"
-version = "1.21.0"
+version = "1.23.1"
source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "ef657dfab802224e671f5818e9a4935f9b1957ed18e58292690cc39e7a4092a3"
+checksum = "5c76a5792e44e4abe34d3abf15636779261d45a7450612059293d1d2cfc63422"
[[package]]
name = "byteorder"
@@ -619,22 +619,30 @@ dependencies = [
"libc",
]
+[[package]]
+name = "bzip2"
+version = "0.5.2"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "49ecfb22d906f800d4fe833b6282cf4dc1c298f5057ca0b5445e5c209735ca47"
+dependencies = [
+ "bzip2-sys",
+]
+
[[package]]
name = "bzip2-sys"
-version = "0.1.11+1.0.8"
+version = "0.1.13+1.0.8"
source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "736a955f3fa7875102d57c82b8cac37ec45224a07fd32d58f9f7a186b6cd4cdc"
+checksum = "225bff33b2141874fe80d71e07d6eec4f85c5c216453dd96388240f96e1acc14"
dependencies = [
"cc",
- "libc",
"pkg-config",
]
[[package]]
name = "cc"
-version = "1.2.29"
+version = "1.2.31"
source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "5c1599538de2394445747c8cf7935946e3cc27e9625f889d979bfb2aaf569362"
+checksum = "c3a42d84bb6b69d3a8b3eaacf0d88f179e1929695e1ad012b6cf64d9caaa5fd2"
dependencies = [
"jobserver",
"libc",
@@ -658,9 +666,9 @@ dependencies = [
[[package]]
name = "cfg-if"
-version = "1.0.0"
+version = "1.0.1"
source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "baf1de4339761588bc0619e3cbc0120ee582ebb74b53b4efbf79117bd2da40fd"
+checksum = "9555578bc9e57714c812a1f84e4fc5b4d21fcb063490c624de019f7464c91268"
[[package]]
name = "cfg_aliases"
@@ -679,7 +687,7 @@ dependencies = [
"js-sys",
"num-traits",
"wasm-bindgen",
- "windows-targets",
+ "windows-targets 0.52.6",
]
[[package]]
@@ -736,12 +744,11 @@ dependencies = [
[[package]]
name = "comfy-table"
-version = "7.1.3"
+version = "7.1.4"
source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "24f165e7b643266ea80cb858aed492ad9280e3e05ce24d4a99d7d7b889b6a4d9"
+checksum = "4a65ebfec4fb190b6f90e944a817d60499ee0744e582530e2c9900a22e591d9a"
dependencies = [
- "strum",
- "strum_macros",
+ "unicode-segmentation",
"unicode-width",
]
@@ -760,7 +767,7 @@ version = "0.1.16"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "f9d839f2a20b0aee515dc581a6172f2321f96cab76c1a38a4c584a194955390e"
dependencies = [
- "getrandom 0.2.15",
+ "getrandom 0.2.16",
"once_cell",
"tiny-keccak",
]
@@ -794,18 +801,18 @@ dependencies = [
[[package]]
name = "cpufeatures"
-version = "0.2.16"
+version = "0.2.17"
source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "16b80225097f2e5ae4e7179dd2266824648f3e2f49d9134d584b76389d31c4c3"
+checksum = "59ed5838eebb26a2bb2e58f6d5b5316989ae9d08bab10e0e6d103e656d1b0280"
dependencies = [
"libc",
]
[[package]]
name = "crc32fast"
-version = "1.4.2"
+version = "1.5.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "a97769d94ddab943e4510d138150169a2758b5ef3eb191a9ee688de3e23ef7b3"
+checksum = "9481c1c90cbf2ac953f07c8d4a58aa3945c425b7185c9154d67a65e4230da511"
dependencies = [
"cfg-if",
]
@@ -818,9 +825,9 @@ checksum = "d0a5c400df2834b80a4c3327b3aad3a4c4cd4de0629063962b03235697506a28"
[[package]]
name = "crunchy"
-version = "0.2.2"
+version = "0.2.4"
source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "7a81dae078cea95a014a339291cec439d2f232ebe854a9d672b796c6afafa9b7"
+checksum = "460fbee9c2c2f33933d720630a6a0bac33ba7053db5344fac858d4b8952d77d5"
[[package]]
name = "crypto-common"
@@ -846,9 +853,9 @@ dependencies = [
[[package]]
name = "csv-core"
-version = "0.1.11"
+version = "0.1.12"
source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "5efa2b3d7902f4b634a20cae3c9c4e6209dc4779feb6863329607560143efa70"
+checksum = "7d02f3b0da4c6504f86e9cd789d8dbafab48c2321be74e9987593de5a894d93d"
dependencies = [
"memchr",
]
@@ -870,7 +877,7 @@ dependencies = [
[[package]]
name = "datafusion"
version = "42.0.0"
-source = "git+https://github.com/blaze-init/arrow-datafusion.git?rev=b7876646a02401cc064657ccf17de27cfe6240ee#b7876646a02401cc064657ccf17de27cfe6240ee"
+source = "git+https://github.com/blaze-init/arrow-datafusion.git?rev=0e6a0b0942cd052d8c2a3606f26d9823afc911d9#0e6a0b0942cd052d8c2a3606f26d9823afc911d9"
dependencies = [
"ahash",
"arrow",
@@ -880,7 +887,7 @@ dependencies = [
"async-compression",
"async-trait",
"bytes 1.10.1",
- "bzip2",
+ "bzip2 0.4.4",
"chrono",
"dashmap",
"datafusion-catalog",
@@ -926,7 +933,7 @@ dependencies = [
[[package]]
name = "datafusion-catalog"
version = "42.0.0"
-source = "git+https://github.com/blaze-init/arrow-datafusion.git?rev=b7876646a02401cc064657ccf17de27cfe6240ee#b7876646a02401cc064657ccf17de27cfe6240ee"
+source = "git+https://github.com/blaze-init/arrow-datafusion.git?rev=0e6a0b0942cd052d8c2a3606f26d9823afc911d9#0e6a0b0942cd052d8c2a3606f26d9823afc911d9"
dependencies = [
"arrow-schema",
"async-trait",
@@ -940,7 +947,7 @@ dependencies = [
[[package]]
name = "datafusion-common"
version = "42.0.0"
-source = "git+https://github.com/blaze-init/arrow-datafusion.git?rev=b7876646a02401cc064657ccf17de27cfe6240ee#b7876646a02401cc064657ccf17de27cfe6240ee"
+source = "git+https://github.com/blaze-init/arrow-datafusion.git?rev=0e6a0b0942cd052d8c2a3606f26d9823afc911d9#0e6a0b0942cd052d8c2a3606f26d9823afc911d9"
dependencies = [
"ahash",
"arrow",
@@ -963,7 +970,7 @@ dependencies = [
[[package]]
name = "datafusion-common-runtime"
version = "42.0.0"
-source = "git+https://github.com/blaze-init/arrow-datafusion.git?rev=b7876646a02401cc064657ccf17de27cfe6240ee#b7876646a02401cc064657ccf17de27cfe6240ee"
+source = "git+https://github.com/blaze-init/arrow-datafusion.git?rev=0e6a0b0942cd052d8c2a3606f26d9823afc911d9#0e6a0b0942cd052d8c2a3606f26d9823afc911d9"
dependencies = [
"log",
"tokio",
@@ -972,7 +979,7 @@ dependencies = [
[[package]]
name = "datafusion-execution"
version = "42.0.0"
-source = "git+https://github.com/blaze-init/arrow-datafusion.git?rev=b7876646a02401cc064657ccf17de27cfe6240ee#b7876646a02401cc064657ccf17de27cfe6240ee"
+source = "git+https://github.com/blaze-init/arrow-datafusion.git?rev=0e6a0b0942cd052d8c2a3606f26d9823afc911d9#0e6a0b0942cd052d8c2a3606f26d9823afc911d9"
dependencies = [
"arrow",
"chrono",
@@ -992,7 +999,7 @@ dependencies = [
[[package]]
name = "datafusion-expr"
version = "42.0.0"
-source = "git+https://github.com/blaze-init/arrow-datafusion.git?rev=b7876646a02401cc064657ccf17de27cfe6240ee#b7876646a02401cc064657ccf17de27cfe6240ee"
+source = "git+https://github.com/blaze-init/arrow-datafusion.git?rev=0e6a0b0942cd052d8c2a3606f26d9823afc911d9#0e6a0b0942cd052d8c2a3606f26d9823afc911d9"
dependencies = [
"ahash",
"arrow",
@@ -1013,7 +1020,7 @@ dependencies = [
[[package]]
name = "datafusion-expr-common"
version = "42.0.0"
-source = "git+https://github.com/blaze-init/arrow-datafusion.git?rev=b7876646a02401cc064657ccf17de27cfe6240ee#b7876646a02401cc064657ccf17de27cfe6240ee"
+source = "git+https://github.com/blaze-init/arrow-datafusion.git?rev=0e6a0b0942cd052d8c2a3606f26d9823afc911d9#0e6a0b0942cd052d8c2a3606f26d9823afc911d9"
dependencies = [
"arrow",
"datafusion-common",
@@ -1127,7 +1134,7 @@ dependencies = [
[[package]]
name = "datafusion-functions"
version = "42.0.0"
-source = "git+https://github.com/blaze-init/arrow-datafusion.git?rev=b7876646a02401cc064657ccf17de27cfe6240ee#b7876646a02401cc064657ccf17de27cfe6240ee"
+source = "git+https://github.com/blaze-init/arrow-datafusion.git?rev=0e6a0b0942cd052d8c2a3606f26d9823afc911d9#0e6a0b0942cd052d8c2a3606f26d9823afc911d9"
dependencies = [
"arrow",
"arrow-buffer",
@@ -1153,7 +1160,7 @@ dependencies = [
[[package]]
name = "datafusion-functions-aggregate"
version = "42.0.0"
-source = "git+https://github.com/blaze-init/arrow-datafusion.git?rev=b7876646a02401cc064657ccf17de27cfe6240ee#b7876646a02401cc064657ccf17de27cfe6240ee"
+source = "git+https://github.com/blaze-init/arrow-datafusion.git?rev=0e6a0b0942cd052d8c2a3606f26d9823afc911d9#0e6a0b0942cd052d8c2a3606f26d9823afc911d9"
dependencies = [
"ahash",
"arrow",
@@ -1173,7 +1180,7 @@ dependencies = [
[[package]]
name = "datafusion-functions-aggregate-common"
version = "42.0.0"
-source = "git+https://github.com/blaze-init/arrow-datafusion.git?rev=b7876646a02401cc064657ccf17de27cfe6240ee#b7876646a02401cc064657ccf17de27cfe6240ee"
+source = "git+https://github.com/blaze-init/arrow-datafusion.git?rev=0e6a0b0942cd052d8c2a3606f26d9823afc911d9#0e6a0b0942cd052d8c2a3606f26d9823afc911d9"
dependencies = [
"ahash",
"arrow",
@@ -1186,7 +1193,7 @@ dependencies = [
[[package]]
name = "datafusion-functions-nested"
version = "42.0.0"
-source = "git+https://github.com/blaze-init/arrow-datafusion.git?rev=b7876646a02401cc064657ccf17de27cfe6240ee#b7876646a02401cc064657ccf17de27cfe6240ee"
+source = "git+https://github.com/blaze-init/arrow-datafusion.git?rev=0e6a0b0942cd052d8c2a3606f26d9823afc911d9#0e6a0b0942cd052d8c2a3606f26d9823afc911d9"
dependencies = [
"arrow",
"arrow-array",
@@ -1208,7 +1215,7 @@ dependencies = [
[[package]]
name = "datafusion-functions-window"
version = "42.0.0"
-source = "git+https://github.com/blaze-init/arrow-datafusion.git?rev=b7876646a02401cc064657ccf17de27cfe6240ee#b7876646a02401cc064657ccf17de27cfe6240ee"
+source = "git+https://github.com/blaze-init/arrow-datafusion.git?rev=0e6a0b0942cd052d8c2a3606f26d9823afc911d9#0e6a0b0942cd052d8c2a3606f26d9823afc911d9"
dependencies = [
"datafusion-common",
"datafusion-expr",
@@ -1219,7 +1226,7 @@ dependencies = [
[[package]]
name = "datafusion-optimizer"
version = "42.0.0"
-source = "git+https://github.com/blaze-init/arrow-datafusion.git?rev=b7876646a02401cc064657ccf17de27cfe6240ee#b7876646a02401cc064657ccf17de27cfe6240ee"
+source = "git+https://github.com/blaze-init/arrow-datafusion.git?rev=0e6a0b0942cd052d8c2a3606f26d9823afc911d9#0e6a0b0942cd052d8c2a3606f26d9823afc911d9"
dependencies = [
"arrow",
"async-trait",
@@ -1238,7 +1245,7 @@ dependencies = [
[[package]]
name = "datafusion-physical-expr"
version = "42.0.0"
-source = "git+https://github.com/blaze-init/arrow-datafusion.git?rev=b7876646a02401cc064657ccf17de27cfe6240ee#b7876646a02401cc064657ccf17de27cfe6240ee"
+source = "git+https://github.com/blaze-init/arrow-datafusion.git?rev=0e6a0b0942cd052d8c2a3606f26d9823afc911d9#0e6a0b0942cd052d8c2a3606f26d9823afc911d9"
dependencies = [
"ahash",
"arrow",
@@ -1262,14 +1269,14 @@ dependencies = [
"itertools 0.13.0",
"log",
"paste",
- "petgraph",
+ "petgraph 0.6.5",
"regex",
]
[[package]]
name = "datafusion-physical-expr-common"
version = "42.0.0"
-source = "git+https://github.com/blaze-init/arrow-datafusion.git?rev=b7876646a02401cc064657ccf17de27cfe6240ee#b7876646a02401cc064657ccf17de27cfe6240ee"
+source = "git+https://github.com/blaze-init/arrow-datafusion.git?rev=0e6a0b0942cd052d8c2a3606f26d9823afc911d9#0e6a0b0942cd052d8c2a3606f26d9823afc911d9"
dependencies = [
"ahash",
"arrow",
@@ -1282,7 +1289,7 @@ dependencies = [
[[package]]
name = "datafusion-physical-optimizer"
version = "42.0.0"
-source = "git+https://github.com/blaze-init/arrow-datafusion.git?rev=b7876646a02401cc064657ccf17de27cfe6240ee#b7876646a02401cc064657ccf17de27cfe6240ee"
+source = "git+https://github.com/blaze-init/arrow-datafusion.git?rev=0e6a0b0942cd052d8c2a3606f26d9823afc911d9#0e6a0b0942cd052d8c2a3606f26d9823afc911d9"
dependencies = [
"arrow-schema",
"datafusion-common",
@@ -1295,7 +1302,7 @@ dependencies = [
[[package]]
name = "datafusion-physical-plan"
version = "42.0.0"
-source = "git+https://github.com/blaze-init/arrow-datafusion.git?rev=b7876646a02401cc064657ccf17de27cfe6240ee#b7876646a02401cc064657ccf17de27cfe6240ee"
+source = "git+https://github.com/blaze-init/arrow-datafusion.git?rev=0e6a0b0942cd052d8c2a3606f26d9823afc911d9#0e6a0b0942cd052d8c2a3606f26d9823afc911d9"
dependencies = [
"ahash",
"arrow",
@@ -1329,7 +1336,7 @@ dependencies = [
[[package]]
name = "datafusion-sql"
version = "42.0.0"
-source = "git+https://github.com/blaze-init/arrow-datafusion.git?rev=b7876646a02401cc064657ccf17de27cfe6240ee#b7876646a02401cc064657ccf17de27cfe6240ee"
+source = "git+https://github.com/blaze-init/arrow-datafusion.git?rev=0e6a0b0942cd052d8c2a3606f26d9823afc911d9#0e6a0b0942cd052d8c2a3606f26d9823afc911d9"
dependencies = [
"arrow",
"arrow-array",
@@ -1370,7 +1377,7 @@ checksum = "97369cbbc041bc366949bc74d34658d6cda5621039731c6310521892a3a20ae0"
dependencies = [
"proc-macro2",
"quote",
- "syn 2.0.91",
+ "syn 2.0.104",
]
[[package]]
@@ -1381,44 +1388,44 @@ checksum = "92773504d58c093f6de2459af4af33faa518c13451eb8f2b5698ed3d36e7c813"
[[package]]
name = "either"
-version = "1.13.0"
+version = "1.15.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "60b1af1c220855b6ceac025d3f6ecdd2b7c4894bfe9cd9bda4fbb4bc7c0d4cf0"
+checksum = "48c757948c5ede0e46177b7add2e67155f70e33c07fea8284df6576da70b3719"
[[package]]
name = "equator"
-version = "0.2.2"
+version = "0.4.2"
source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "c35da53b5a021d2484a7cc49b2ac7f2d840f8236a286f84202369bd338d761ea"
+checksum = "4711b213838dfee0117e3be6ac926007d7f433d7bbe33595975d4190cb07e6fc"
dependencies = [
"equator-macro",
]
[[package]]
name = "equator-macro"
-version = "0.2.1"
+version = "0.4.2"
source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "3bf679796c0322556351f287a51b49e48f7c4986e727b5dd78c972d30e2e16cc"
+checksum = "44f23cf4b44bfce11a86ace86f8a73ffdec849c9fd00a386a53d278bd9e81fb3"
dependencies = [
"proc-macro2",
"quote",
- "syn 2.0.91",
+ "syn 2.0.104",
]
[[package]]
name = "equivalent"
-version = "1.0.1"
+version = "1.0.2"
source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "5443807d6dff69373d433ab9ef5378ad8df50ca6298caf15de6e52e24aaf54d5"
+checksum = "877a4ace8713b0bcf2a4e7eec82529c029f1d0619886d18145fea96c3ffe5c0f"
[[package]]
name = "errno"
-version = "0.3.10"
+version = "0.3.13"
source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "33d852cb9b869c2a9b3df2f71a3074817f01e1844f839a144f5fcef059a4eb5d"
+checksum = "778e2ac28f6c47af28e4907f13ffd1e1ddbd400980a9abd7c8df189bf578a5ad"
dependencies = [
"libc",
- "windows-sys 0.59.0",
+ "windows-sys 0.60.2",
]
[[package]]
@@ -1435,9 +1442,9 @@ checksum = "37909eebbb50d72f9059c3b6d82c0463f2ff062c9e95845c43a6c9c0355411be"
[[package]]
name = "faststr"
-version = "0.2.27"
+version = "0.2.31"
source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "9154486833a83cb5d99de8c4d831314b8ae810dd4ef18d89ceb7a9c7c728dd74"
+checksum = "a6503af7917fea18ffef8f7e8553fb8dff89e2e6837e94e09dd7fb069c82d62c"
dependencies = [
"bytes 1.10.1",
"rkyv",
@@ -1463,6 +1470,12 @@ version = "0.4.2"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "0ce7134b9999ecaf8bcd65542e436736ef32ddca1b3e06094cb6ec5755203b80"
+[[package]]
+name = "fixedbitset"
+version = "0.5.7"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "1d674e81391d1e1ab681a28d99df07927c6d4aa5b027d7da16ba32d1d21ecd99"
+
[[package]]
name = "flatbuffers"
version = "24.12.23"
@@ -1475,9 +1488,9 @@ dependencies = [
[[package]]
name = "flate2"
-version = "1.0.35"
+version = "1.1.2"
source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "c936bfdafb507ebbf50b8074c54fa31c5be9a1e7e5f467dd659697041407d07c"
+checksum = "4a3d7db9596fecd151c5f638c0ee5d5bd487b6e0ea232e5dc96d5250f6f94b1d"
dependencies = [
"crc32fast",
"miniz_oxide",
@@ -1572,7 +1585,7 @@ checksum = "162ee34ebcb7c64a8abebc059ce0fee27c2262618d7b60ed8faf72fef13c3650"
dependencies = [
"proc-macro2",
"quote",
- "syn 2.0.91",
+ "syn 2.0.104",
]
[[package]]
@@ -1629,25 +1642,25 @@ dependencies = [
[[package]]
name = "getrandom"
-version = "0.2.15"
+version = "0.2.16"
source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "c4567c8db10ae91089c99af84c68c38da3ec2f087c3f82960bcdbf3656b6f4d7"
+checksum = "335ff9f135e4384c8150d6f27c6daed433577f86b4750418338c01a1a2528592"
dependencies = [
"cfg-if",
"libc",
- "wasi 0.11.0+wasi-snapshot-preview1",
+ "wasi 0.11.1+wasi-snapshot-preview1",
]
[[package]]
name = "getrandom"
-version = "0.3.1"
+version = "0.3.3"
source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "43a49c392881ce6d5c3b8cb70f98717b7c07aabbdff06687b9030dbfbe2725f8"
+checksum = "26145e563e54f2cadc477553f1ec5ee650b00862f0a58bcd12cbdc5f0ea2d2f4"
dependencies = [
"cfg-if",
"libc",
- "wasi 0.13.3+wasi-0.2.2",
- "windows-targets",
+ "r-efi",
+ "wasi 0.14.2+wasi-0.2.4",
]
[[package]]
@@ -1658,9 +1671,9 @@ checksum = "07e28edb80900c19c28f1072f2e8aeca7fa06b23cd4169cefe1af5aa3260783f"
[[package]]
name = "glob"
-version = "0.3.1"
+version = "0.3.2"
source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "d2fabcfbdc87f4758337ca535fb41a6d701b65693ce38287d856d1674551ec9b"
+checksum = "a8d1add55171497b4705a648c6b583acafb01d58050a51727785f0b2c8e0a2b2"
[[package]]
name = "h2"
@@ -1683,9 +1696,9 @@ dependencies = [
[[package]]
name = "half"
-version = "2.4.1"
+version = "2.6.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "6dd08c532ae367adf81c312a4580bc67f1d0fe8bc9c460520283f4c0ff277888"
+checksum = "459196ed295495a68f7d7fe1d84f6c4b7ff0e21fe3017b2f283c6fac3ad803c9"
dependencies = [
"cfg-if",
"crunchy",
@@ -1704,9 +1717,9 @@ dependencies = [
[[package]]
name = "hashbrown"
-version = "0.15.2"
+version = "0.15.4"
source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "bf151400ff0baff5465007dd2f3e717f3fe502074ca563069ce3a6629d07b289"
+checksum = "5971ac85611da7067dbfcabef3c70ebb5606018acd9e2a3903a0da507521e0d5"
[[package]]
name = "headers"
@@ -1740,15 +1753,9 @@ checksum = "2304e00983f87ffb38b55b444b5e3b60a884b5d30c0fca7d82fe33449bbe55ea"
[[package]]
name = "hermit-abi"
-version = "0.3.9"
-source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "d231dfb89cfffdbc30e7fc41579ed6066ad03abda9e567ccafae602b97ec5024"
-
-[[package]]
-name = "hermit-abi"
-version = "0.4.0"
+version = "0.5.2"
source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "fbf6a919d6cf397374f7dfeeea91d974c7c0a7221d0d0f4f20d859d329e53fcc"
+checksum = "fc0fef456e4baa96da950455cd02c081ca953b141298e41db3fc7e36b1da849c"
[[package]]
name = "hex"
@@ -1813,9 +1820,9 @@ checksum = "df3b46402a9d5adb4c86a0cf463f42e19994e3ee891101b1841f30a545cb49a9"
[[package]]
name = "humantime"
-version = "2.1.0"
+version = "2.2.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "9a3a5bfb195931eeb336b2a7b4d761daec841b97f947d34394601737a7bba5e4"
+checksum = "9b112acc8b3adf4b107a8ec20977da0273a8c386765a3ec0229bd500a1443f9f"
[[package]]
name = "hyper"
@@ -1833,7 +1840,7 @@ dependencies = [
"httpdate",
"itoa",
"pin-project-lite",
- "smallvec 1.13.2",
+ "smallvec 1.15.1",
"tokio",
]
@@ -1854,14 +1861,15 @@ dependencies = [
[[package]]
name = "iana-time-zone"
-version = "0.1.61"
+version = "0.1.63"
source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "235e081f3925a06703c2d0117ea8b91f042756fd6e7a6e5d901e8ca1a996b220"
+checksum = "b0c919e5debc312ad217002b8048a17b7d83f80703865bbfcfebb0458b0b27d8"
dependencies = [
"android_system_properties",
"core-foundation-sys",
"iana-time-zone-haiku",
"js-sys",
+ "log",
"wasm-bindgen",
"windows-core",
]
@@ -1877,21 +1885,22 @@ dependencies = [
[[package]]
name = "icu_collections"
-version = "1.5.0"
+version = "2.0.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "db2fa452206ebee18c4b5c2274dbf1de17008e874b4dc4f0aea9d01ca79e4526"
+checksum = "200072f5d0e3614556f94a9930d5dc3e0662a652823904c3a75dc3b0af7fee47"
dependencies = [
"displaydoc",
+ "potential_utf",
"yoke",
"zerofrom",
"zerovec",
]
[[package]]
-name = "icu_locid"
-version = "1.5.0"
+name = "icu_locale_core"
+version = "2.0.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "13acbb8371917fc971be86fc8057c41a64b521c184808a698c02acc242dbf637"
+checksum = "0cde2700ccaed3872079a65fb1a78f6c0a36c91570f28755dda67bc8f7d9f00a"
dependencies = [
"displaydoc",
"litemap",
@@ -1900,99 +1909,66 @@ dependencies = [
"zerovec",
]
-[[package]]
-name = "icu_locid_transform"
-version = "1.5.0"
-source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "01d11ac35de8e40fdeda00d9e1e9d92525f3f9d887cdd7aa81d727596788b54e"
-dependencies = [
- "displaydoc",
- "icu_locid",
- "icu_locid_transform_data",
- "icu_provider",
- "tinystr",
- "zerovec",
-]
-
-[[package]]
-name = "icu_locid_transform_data"
-version = "1.5.0"
-source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "fdc8ff3388f852bede6b579ad4e978ab004f139284d7b28715f773507b946f6e"
-
[[package]]
name = "icu_normalizer"
-version = "1.5.0"
+version = "2.0.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "19ce3e0da2ec68599d193c93d088142efd7f9c5d6fc9b803774855747dc6a84f"
+checksum = "436880e8e18df4d7bbc06d58432329d6458cc84531f7ac5f024e93deadb37979"
dependencies = [
"displaydoc",
"icu_collections",
"icu_normalizer_data",
"icu_properties",
"icu_provider",
- "smallvec 1.13.2",
- "utf16_iter",
- "utf8_iter",
- "write16",
+ "smallvec 1.15.1",
"zerovec",
]
[[package]]
name = "icu_normalizer_data"
-version = "1.5.0"
+version = "2.0.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "f8cafbf7aa791e9b22bec55a167906f9e1215fd475cd22adfcf660e03e989516"
+checksum = "00210d6893afc98edb752b664b8890f0ef174c8adbb8d0be9710fa66fbbf72d3"
[[package]]
name = "icu_properties"
-version = "1.5.1"
+version = "2.0.1"
source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "93d6020766cfc6302c15dbbc9c8778c37e62c14427cb7f6e601d849e092aeef5"
+checksum = "016c619c1eeb94efb86809b015c58f479963de65bdb6253345c1a1276f22e32b"
dependencies = [
"displaydoc",
"icu_collections",
- "icu_locid_transform",
+ "icu_locale_core",
"icu_properties_data",
"icu_provider",
- "tinystr",
+ "potential_utf",
+ "zerotrie",
"zerovec",
]
[[package]]
name = "icu_properties_data"
-version = "1.5.0"
+version = "2.0.1"
source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "67a8effbc3dd3e4ba1afa8ad918d5684b8868b3b26500753effea8d2eed19569"
+checksum = "298459143998310acd25ffe6810ed544932242d3f07083eee1084d83a71bd632"
[[package]]
name = "icu_provider"
-version = "1.5.0"
+version = "2.0.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "6ed421c8a8ef78d3e2dbc98a973be2f3770cb42b606e3ab18d6237c4dfde68d9"
+checksum = "03c80da27b5f4187909049ee2d72f276f0d9f99a42c306bd0131ecfe04d8e5af"
dependencies = [
"displaydoc",
- "icu_locid",
- "icu_provider_macros",
+ "icu_locale_core",
"stable_deref_trait",
"tinystr",
"writeable",
"yoke",
"zerofrom",
+ "zerotrie",
"zerovec",
]
-[[package]]
-name = "icu_provider_macros"
-version = "1.5.0"
-source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "1ec89e9337638ecdc08744df490b221a7399bf8d164eb52a665454e60e075ad6"
-dependencies = [
- "proc-macro2",
- "quote",
- "syn 2.0.91",
-]
-
[[package]]
name = "idna"
version = "1.0.3"
@@ -2000,15 +1976,15 @@ source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "686f825264d630750a544639377bae737628043f20d38bbc029e8f29ea968a7e"
dependencies = [
"idna_adapter",
- "smallvec 1.13.2",
+ "smallvec 1.15.1",
"utf8_iter",
]
[[package]]
name = "idna_adapter"
-version = "1.2.0"
+version = "1.2.1"
source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "daca1df1c957320b2cf139ac61e7bd64fed304c5040df000a745aa1de3b4ef71"
+checksum = "3acae9609540aa318d1bc588455225fb2085b9ed0c4f6bd0d9d5bcd86f1a0344"
dependencies = [
"icu_normalizer",
"icu_properties",
@@ -2016,12 +1992,12 @@ dependencies = [
[[package]]
name = "indexmap"
-version = "2.7.0"
+version = "2.10.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "62f822373a4fe84d4bb149bf54e584a7f4abec90e072ed49cda0edea5b95471f"
+checksum = "fe4cd85333e22411419a0bcae1297d25e58c9443848b11dc6a86fefe8c78a661"
dependencies = [
"equivalent",
- "hashbrown 0.15.2",
+ "hashbrown 0.15.4",
]
[[package]]
@@ -2062,22 +2038,22 @@ checksum = "8bb03732005da905c88227371639bf1ad885cc712789c011c31c5fb3ab3ccf02"
[[package]]
name = "io-uring"
-version = "0.7.8"
+version = "0.7.9"
source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "b86e202f00093dcba4275d4636b93ef9dd75d025ae560d2521b45ea28ab49013"
+checksum = "d93587f37623a1a17d94ef2bc9ada592f5465fe7732084ab7beefabe5c77c0c4"
dependencies = [
- "bitflags 2.6.0",
+ "bitflags 2.9.1",
"cfg-if",
"libc",
]
[[package]]
name = "is-terminal"
-version = "0.4.15"
+version = "0.4.16"
source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "e19b23d53f35ce9f56aebc7d1bb4e6ac1e9c0db7ac85c8d1760c04379edced37"
+checksum = "e04d7f318608d35d4b61ddd75cbdaee86b023ebe2bd5a66ee0915f0bf93095a9"
dependencies = [
- "hermit-abi 0.4.0",
+ "hermit-abi",
"libc",
"windows-sys 0.59.0",
]
@@ -2111,15 +2087,15 @@ dependencies = [
[[package]]
name = "itoa"
-version = "1.0.14"
+version = "1.0.15"
source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "d75a2a4b1b190afb6f5425f10f6a8f959d2ea0b9c2b1d79553551850539e4674"
+checksum = "4a5f13b858c8d314ee3e8f639011f7ccefe71f97f96e50151fb991f267928e2c"
[[package]]
name = "jemalloc_pprof"
-version = "0.8.0"
+version = "0.8.1"
source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "9770dc31e3b7797ce9ec1184c23021eae0ee50d8cac2f3c706a9cae1fb664e41"
+checksum = "74ff642505c7ce8d31c0d43ec0e235c6fd4585d9b8172d8f9dd04d36590200b5"
dependencies = [
"anyhow",
"libc",
@@ -2154,10 +2130,11 @@ checksum = "8eaf4bc02d17cbdd7ff4c7438cafcdf7fb9a4613313ad11b4f8fefe7d3fa0130"
[[package]]
name = "jobserver"
-version = "0.1.32"
+version = "0.1.33"
source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "48d1dbcbbeb6a7fec7e059840aa538bd62aaccf972c7346c4d9d2059312853d0"
+checksum = "38f262f097c174adebe41eb73d66ae9c06b2844fb0da69969647bbddd9b0538a"
dependencies = [
+ "getrandom 0.3.3",
"libc",
]
@@ -2249,9 +2226,9 @@ dependencies = [
[[package]]
name = "libc"
-version = "0.2.172"
+version = "0.2.174"
source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "d750af042f7ef4f724306de029d18836c26c1765a54a6a3f094cbd23a7267ffa"
+checksum = "1171693293099992e19cddea4e8b849964e9846f4acee11b3948bcc337be8776"
[[package]]
name = "libloading"
@@ -2260,14 +2237,14 @@ source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "07033963ba89ebaf1584d767badaa2e8fcec21aedea6b8c0346d487d49c28667"
dependencies = [
"cfg-if",
- "windows-targets",
+ "windows-targets 0.53.3",
]
[[package]]
name = "libm"
-version = "0.2.11"
+version = "0.2.15"
source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "8355be11b20d696c8f18f6cc018c4e372165b1fa8126cef092399c9951984ffa"
+checksum = "f9fbbcab51052fe104eb5e5d351cf728d30a5be1fe14d9be8a3b097481fb97de"
[[package]]
name = "linux-raw-sys"
@@ -2277,15 +2254,15 @@ checksum = "d26c52dbd32dccf2d10cac7725f8eae5296885fb5703b261f7d0a0739ec807ab"
[[package]]
name = "linux-raw-sys"
-version = "0.9.3"
+version = "0.9.4"
source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "fe7db12097d22ec582439daf8618b8fdd1a7bef6270e9af3b1ebcd30893cf413"
+checksum = "cd945864f07fe9f5371a27ad7b52a172b4b499999f1d97574c9fa68373937e12"
[[package]]
name = "litemap"
-version = "0.7.4"
+version = "0.8.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "4ee93343901ab17bd981295f2cf0026d4ad018c7c31ba84549a4ddbb47a45104"
+checksum = "241eaef5fd12c88705a01fc1066c48c4b36e0dd4377dcdc7ec3942cea7a69956"
[[package]]
name = "lock_api"
@@ -2335,9 +2312,9 @@ dependencies = [
[[package]]
name = "mappings"
-version = "0.7.0"
+version = "0.7.1"
source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "e434981a332777c2b3062652d16a55f8e74fa78e6b1882633f0d77399c84fc2a"
+checksum = "db4d277bb50d4508057e7bddd7fcd19ef4a4cc38051b6a5a36868d75ae2cbeb9"
dependencies = [
"anyhow",
"libc",
@@ -2358,15 +2335,15 @@ dependencies = [
[[package]]
name = "memchr"
-version = "2.7.4"
+version = "2.7.5"
source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "78ca9ab1a0babb1e7d5695e3530886289c18cf2f87ec19a575a0abdce112e3a3"
+checksum = "32a282da65faaf38286cf3be983213fcf1d2e2a58700e808f83f4ea9a4804bc0"
[[package]]
name = "memmap2"
-version = "0.9.5"
+version = "0.9.7"
source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "fd3f7eed9d3848f8b98834af67102b720745c4ec028fcd0aa0239277e7de374f"
+checksum = "483758ad303d734cec05e5c12b41d7e93e6a6390c5e9dae6bdeb7c1259012d28"
dependencies = [
"libc",
]
@@ -2385,9 +2362,9 @@ checksum = "68354c5c6bd36d73ff3feceb05efa59b6acb7626617f4962be322a825e61f79a"
[[package]]
name = "miniz_oxide"
-version = "0.8.2"
+version = "0.8.9"
source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "4ffbe83022cedc1d264172192511ae958937694cd57ce297164951b8b3568394"
+checksum = "1fa76a2c86f704bdb222d66965fb3d63269ce38518b83cb0575fca855ebb6316"
dependencies = [
"adler2",
]
@@ -2399,7 +2376,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "78bed444cc8a2160f01cbcf811ef18cac863ad68ae8ca62092e8db51d51c761c"
dependencies = [
"libc",
- "wasi 0.11.0+wasi-snapshot-preview1",
+ "wasi 0.11.1+wasi-snapshot-preview1",
"windows-sys 0.59.0",
]
@@ -2411,22 +2388,22 @@ checksum = "1d87ecb2933e8aeadb3e3a02b828fed80a7528047e68b4f424523a0981a3a084"
[[package]]
name = "munge"
-version = "0.4.1"
+version = "0.4.6"
source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "64142d38c84badf60abf06ff9bd80ad2174306a5b11bd4706535090a30a419df"
+checksum = "d7feb0b48aa0a25f9fe0899482c6e1379ee7a11b24a53073eacdecb9adb6dc60"
dependencies = [
"munge_macro",
]
[[package]]
name = "munge_macro"
-version = "0.4.1"
+version = "0.4.6"
source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "1bb5c1d8184f13f7d0ccbeeca0def2f9a181bce2624302793005f5ca8aa62e5e"
+checksum = "f2e3795a5d2da581a8b252fec6022eee01aea10161a4d1bf237d4cbe47f7e988"
dependencies = [
"proc-macro2",
"quote",
- "syn 2.0.91",
+ "syn 2.0.104",
]
[[package]]
@@ -2446,7 +2423,7 @@ version = "0.30.1"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "74523f3a35e05aba87a1d978330aef40f67b0304ac79c1c00b294c9830543db6"
dependencies = [
- "bitflags 2.6.0",
+ "bitflags 2.9.1",
"cfg-if",
"cfg_aliases",
"libc",
@@ -2548,11 +2525,11 @@ dependencies = [
[[package]]
name = "num_cpus"
-version = "1.16.0"
+version = "1.17.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "4161fcb6d602d4d2081af7c3a45852d875a03dd337a6bfdd6e06407b61342a43"
+checksum = "91df4bbde75afed763b708b7eee1e8e7651e02d97f6d5dd763e89367e957b23b"
dependencies = [
- "hermit-abi 0.3.9",
+ "hermit-abi",
"libc",
]
@@ -2567,9 +2544,9 @@ dependencies = [
[[package]]
name = "object_store"
-version = "0.11.1"
+version = "0.11.2"
source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "6eb4c22c6154a1e759d7099f9ffad7cc5ef8245f9efbab4a41b92623079c82f3"
+checksum = "3cfccb68961a56facde1163f9319e0d15743352344e7808a11795fb99698dcaf"
dependencies = [
"async-trait",
"bytes 1.10.1",
@@ -2595,7 +2572,7 @@ checksum = "42f5e15c9953c5e4ccceeb2e7382a716482c34515315f7b03532b8b4e8393d2d"
[[package]]
name = "orc-rust"
version = "0.4.1"
-source = "git+https://github.com/blaze-init/datafusion-orc.git?rev=172183facb1d05d6ac0b4ec234207a52387db2a3#172183facb1d05d6ac0b4ec234207a52387db2a3"
+source = "git+https://github.com/blaze-init/datafusion-orc.git?rev=6499dffa74267623ee311b1ca68db34237e83b00#6499dffa74267623ee311b1ca68db34237e83b00"
dependencies = [
"arrow",
"async-trait",
@@ -2651,14 +2628,14 @@ dependencies = [
"cfg-if",
"libc",
"redox_syscall",
- "smallvec 1.13.2",
- "windows-targets",
+ "smallvec 1.15.1",
+ "windows-targets 0.52.6",
]
[[package]]
name = "parquet"
version = "53.0.0"
-source = "git+https://github.com/blaze-init/arrow-rs.git?rev=1fbdc9f980aa34c7758da9db79a64248f24daf1e#1fbdc9f980aa34c7758da9db79a64248f24daf1e"
+source = "git+https://github.com/blaze-init/arrow-rs.git?rev=94ef6f608254484c0ecfc33fcef19891c4f87612#94ef6f608254484c0ecfc33fcef19891c4f87612"
dependencies = [
"ahash",
"arrow-array",
@@ -2717,24 +2694,34 @@ version = "0.6.5"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "b4c5cc86750666a3ed20bdaf5ca2a0344f9c67674cae0515bec2da16fbaa47db"
dependencies = [
- "fixedbitset",
+ "fixedbitset 0.4.2",
+ "indexmap",
+]
+
+[[package]]
+name = "petgraph"
+version = "0.7.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "3672b37090dbd86368a4145bc067582552b29c27377cad4e0a306c97f9bd7772"
+dependencies = [
+ "fixedbitset 0.5.7",
"indexmap",
]
[[package]]
name = "phf"
-version = "0.11.2"
+version = "0.11.3"
source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "ade2d8b8f33c7333b51bcf0428d37e217e9f32192ae4772156f65063b8ce03dc"
+checksum = "1fd6780a80ae0c52cc120a26a1a42c1ae51b247a253e4e06113d23d2c2edd078"
dependencies = [
"phf_shared",
]
[[package]]
name = "phf_codegen"
-version = "0.11.2"
+version = "0.11.3"
source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "e8d39688d359e6b34654d328e262234662d16cc0f60ec8dcbe5e718709342a5a"
+checksum = "aef8048c789fa5e851558d709946d6d79a8ff88c0440c587967f8e94bfb1216a"
dependencies = [
"phf_generator",
"phf_shared",
@@ -2742,9 +2729,9 @@ dependencies = [
[[package]]
name = "phf_generator"
-version = "0.11.2"
+version = "0.11.3"
source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "48e4cc64c2ad9ebe670cb8fd69dd50ae301650392e81c05f9bfcb2d5bdbc24b0"
+checksum = "3c80231409c20246a13fddb31776fb942c38553c51e871f8cbd687a4cfb5843d"
dependencies = [
"phf_shared",
"rand 0.8.5",
@@ -2752,9 +2739,9 @@ dependencies = [
[[package]]
name = "phf_shared"
-version = "0.11.2"
+version = "0.11.3"
source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "90fcb95eef784c2ac79119d1dd819e162b5da872ce6f3c3abe1e8ca1c082f72b"
+checksum = "67eabc2ef2a60eb7faa00097bd1ffdb5bd28e62bf39990626a582201b7a754e5"
dependencies = [
"siphasher",
]
@@ -2781,9 +2768,9 @@ dependencies = [
[[package]]
name = "pin-project-lite"
-version = "0.2.15"
+version = "0.2.16"
source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "915a1e146535de9163f3987b8944ed8cf49a18bb0056bcebcdcece385cece4ff"
+checksum = "3b3cff922bd51709b605d9ead9aa71031d81447142d828eb4a6eba76fe619f9b"
[[package]]
name = "pin-utils"
@@ -2793,9 +2780,9 @@ checksum = "8b870d8c151b6f2fb93e84a13146138f05d02ed11c7e7c54f8826aaaf7c9f184"
[[package]]
name = "pkg-config"
-version = "0.3.31"
+version = "0.3.32"
source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "953ec861398dccce10c670dfeaf3ec4911ca479e9c02154b3a215178c5f566f2"
+checksum = "7edddbd0b52d732b21ad9a5fab5c704c14cd949e5e9a1ec5929a24fded1b904c"
[[package]]
name = "poem"
@@ -2822,10 +2809,10 @@ dependencies = [
"serde",
"serde_json",
"serde_urlencoded",
- "smallvec 1.13.2",
+ "smallvec 1.15.1",
"sse-codec",
"sync_wrapper",
- "thiserror 2.0.9",
+ "thiserror 2.0.12",
"tokio",
"tokio-rustls",
"tokio-stream",
@@ -2843,7 +2830,16 @@ dependencies = [
"proc-macro-crate",
"proc-macro2",
"quote",
- "syn 2.0.91",
+ "syn 2.0.104",
+]
+
+[[package]]
+name = "potential_utf"
+version = "0.1.2"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "e5a7c30837279ca13e7c867e9e40053bc68740f988cb07f7ca6df43cc734b585"
+dependencies = [
+ "zerovec",
]
[[package]]
@@ -2863,18 +2859,18 @@ dependencies = [
"once_cell",
"protobuf",
"protobuf-codegen",
- "smallvec 1.13.2",
+ "smallvec 1.15.1",
"spin",
"symbolic-demangle",
"tempfile",
- "thiserror 2.0.9",
+ "thiserror 2.0.12",
]
[[package]]
name = "pprof_util"
-version = "0.7.0"
+version = "0.8.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "9fa015c78eed2130951e22c58d2095849391e73817ab2e74f71b0b9f63dd8416"
+checksum = "f9aba4251d95ac86f14c33e688d57a9344bfcff29e9b0c5a063fc66b5facc8a1"
dependencies = [
"anyhow",
"backtrace",
@@ -2886,21 +2882,21 @@ dependencies = [
[[package]]
name = "ppv-lite86"
-version = "0.2.20"
+version = "0.2.21"
source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "77957b295656769bb8ad2b6a6b09d897d94f05c41b069aede1fcdaa675eaea04"
+checksum = "85eae3c4ed2f50dcfe72643da4befc30deadb458a9b590d720cde2f2b1e97da9"
dependencies = [
- "zerocopy 0.7.35",
+ "zerocopy",
]
[[package]]
name = "prettyplease"
-version = "0.2.25"
+version = "0.2.36"
source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "64d1ec885c64d0457d564db4ec299b2dae3f9c02808b8ad9c3a089c591b18033"
+checksum = "ff24dfcda44452b9816fff4cd4227e1bb73ff5a2f1bc1105aa92fb8565ce44d2"
dependencies = [
"proc-macro2",
- "syn 2.0.91",
+ "syn 2.0.104",
]
[[package]]
@@ -2914,9 +2910,9 @@ dependencies = [
[[package]]
name = "proc-macro2"
-version = "1.0.92"
+version = "1.0.95"
source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "37d3544b3f2748c54e147655edb5025752e2303145b5aefb3c3ea2c78b973bb0"
+checksum = "02b3e5e68a3a1a02aad3ec490a98007cbc13c37cbe84a3cd7b8e406d76e7f778"
dependencies = [
"unicode-ident",
]
@@ -2927,7 +2923,7 @@ version = "0.17.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "cc5b72d8145275d844d4b5f6d4e1eef00c8cd889edb6035c21675d1bb1f45c9f"
dependencies = [
- "bitflags 2.6.0",
+ "bitflags 2.9.1",
"chrono",
"flate2",
"hex",
@@ -2941,7 +2937,7 @@ version = "0.17.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "239df02d8349b06fc07398a3a1697b06418223b1c7725085e801e7c0fc6a12ec"
dependencies = [
- "bitflags 2.6.0",
+ "bitflags 2.9.1",
"chrono",
"hex",
]
@@ -2987,12 +2983,12 @@ dependencies = [
"log",
"multimap",
"once_cell",
- "petgraph",
+ "petgraph 0.7.1",
"prettyplease",
"prost 0.13.5",
"prost-types",
"regex",
- "syn 2.0.91",
+ "syn 2.0.104",
"tempfile",
]
@@ -3006,7 +3002,7 @@ dependencies = [
"itertools 0.12.1",
"proc-macro2",
"quote",
- "syn 2.0.91",
+ "syn 2.0.104",
]
[[package]]
@@ -3019,7 +3015,7 @@ dependencies = [
"itertools 0.14.0",
"proc-macro2",
"quote",
- "syn 2.0.91",
+ "syn 2.0.104",
]
[[package]]
@@ -3032,7 +3028,7 @@ dependencies = [
"itertools 0.14.0",
"proc-macro2",
"quote",
- "syn 2.0.91",
+ "syn 2.0.104",
]
[[package]]
@@ -3112,7 +3108,7 @@ checksum = "ca414edb151b4c8d125c12566ab0d74dc9cdba36fb80eb7b848c15f495fd32d1"
dependencies = [
"proc-macro2",
"quote",
- "syn 2.0.91",
+ "syn 2.0.104",
]
[[package]]
@@ -3126,13 +3122,19 @@ dependencies = [
[[package]]
name = "quote"
-version = "1.0.37"
+version = "1.0.40"
source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "b5b9d34b8991d19d98081b46eacdd8eb58c6f2b201139f7c5f643cc155a633af"
+checksum = "1885c039570dc00dcb4ff087a89e185fd56bae234ddc7f056a945bf36467248d"
dependencies = [
"proc-macro2",
]
+[[package]]
+name = "r-efi"
+version = "5.3.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "69cdb34c158ceb288df11e18b4bd39de994f6657d83847bdffdbd7f346754b0f"
+
[[package]]
name = "radium"
version = "0.7.0"
@@ -3166,7 +3168,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "6db2770f06117d490610c7488547d543617b21bfa07796d7a12f6f1bd53850d1"
dependencies = [
"rand_chacha 0.9.0",
- "rand_core 0.9.0",
+ "rand_core 0.9.3",
]
[[package]]
@@ -3186,7 +3188,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "d3022b5f1df60f26e1ffddd6c66e8aa15de382ae63b3a0c1bfc0e4d3e3f325cb"
dependencies = [
"ppv-lite86",
- "rand_core 0.9.0",
+ "rand_core 0.9.3",
]
[[package]]
@@ -3195,46 +3197,45 @@ version = "0.6.4"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "ec0be4795e2f6a28069bec0b5ff3e2ac9bafc99e6a9a7dc3547996c5c816922c"
dependencies = [
- "getrandom 0.2.15",
+ "getrandom 0.2.16",
]
[[package]]
name = "rand_core"
-version = "0.9.0"
+version = "0.9.3"
source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "b08f3c9802962f7e1b25113931d94f43ed9725bebc59db9d0c3e9a23b67e15ff"
+checksum = "99d9a13982dcf210057a8a78572b2217b667c3beacbf3a0d8b454f6f82837d38"
dependencies = [
- "getrandom 0.3.1",
- "zerocopy 0.8.14",
+ "getrandom 0.3.3",
]
[[package]]
name = "redox_syscall"
-version = "0.5.8"
+version = "0.5.17"
source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "03a862b389f93e68874fbf580b9de08dd02facb9a788ebadaf4a3fd33cf58834"
+checksum = "5407465600fb0548f1442edf71dd20683c6ed326200ace4b1ef0763521bb3b77"
dependencies = [
- "bitflags 2.6.0",
+ "bitflags 2.9.1",
]
[[package]]
name = "ref-cast"
-version = "1.0.23"
+version = "1.0.24"
source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "ccf0a6f84d5f1d581da8b41b47ec8600871962f2a528115b542b362d4b744931"
+checksum = "4a0ae411dbe946a674d89546582cea4ba2bb8defac896622d6496f14c23ba5cf"
dependencies = [
"ref-cast-impl",
]
[[package]]
name = "ref-cast-impl"
-version = "1.0.23"
+version = "1.0.24"
source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "bcc303e793d3734489387d205e9b186fac9c6cfacedd98cbb2e8a5943595f3e6"
+checksum = "1165225c21bff1f3bbce98f5a1f889949bc902d3575308cc7b0de30b4f6d27c7"
dependencies = [
"proc-macro2",
"quote",
- "syn 2.0.91",
+ "syn 2.0.104",
]
[[package]]
@@ -3283,9 +3284,9 @@ dependencies = [
[[package]]
name = "rgb"
-version = "0.8.50"
+version = "0.8.52"
source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "57397d16646700483b67d2dd6511d79318f9d057fdbd21a4066aeac8b41d310a"
+checksum = "0c6a884d2998352bb4daf0183589aec883f16a6da1f4dde84d8e2e9a5409a1ce"
dependencies = [
"bytemuck",
]
@@ -3298,7 +3299,7 @@ checksum = "a4689e6c2294d81e88dc6261c768b63bc4fcdb852be6d1352498b114f61383b7"
dependencies = [
"cc",
"cfg-if",
- "getrandom 0.2.15",
+ "getrandom 0.2.16",
"libc",
"untrusted",
"windows-sys 0.52.0",
@@ -3306,12 +3307,12 @@ dependencies = [
[[package]]
name = "rkyv"
-version = "0.8.9"
+version = "0.8.11"
source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "b11a153aec4a6ab60795f8ebe2923c597b16b05bb1504377451e705ef1a45323"
+checksum = "19f5c3e5da784cd8c69d32cdc84673f3204536ca56e1fa01be31a74b92c932ac"
dependencies = [
"bytes 1.10.1",
- "hashbrown 0.15.2",
+ "hashbrown 0.15.4",
"indexmap",
"munge",
"ptr_meta",
@@ -3324,20 +3325,20 @@ dependencies = [
[[package]]
name = "rkyv_derive"
-version = "0.8.9"
+version = "0.8.11"
source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "beb382a4d9f53bd5c0be86b10d8179c3f8a14c30bf774ff77096ed6581e35981"
+checksum = "4270433626cffc9c4c1d3707dd681f2a2718d3d7b09ad754bec137acecda8d22"
dependencies = [
"proc-macro2",
"quote",
- "syn 2.0.91",
+ "syn 2.0.104",
]
[[package]]
name = "rustc-demangle"
-version = "0.1.24"
+version = "0.1.26"
source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "719b953e2095829ee67db738b3bfa9fa368c94900df327b3f07fe6e794d2fe1f"
+checksum = "56f7d92ca342cea22a06f2121d944b4fd82af56988c270852495420f961d4ace"
[[package]]
name = "rustc-hash"
@@ -3360,7 +3361,7 @@ version = "0.38.44"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "fdb5bc1ae2baa591800df16c9ca78619bf65c0488b41b96ccec5d11220d8c154"
dependencies = [
- "bitflags 2.6.0",
+ "bitflags 2.9.1",
"errno",
"libc",
"linux-raw-sys 0.4.15",
@@ -3369,22 +3370,22 @@ dependencies = [
[[package]]
name = "rustix"
-version = "1.0.3"
+version = "1.0.8"
source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "e56a18552996ac8d29ecc3b190b4fdbb2d91ca4ec396de7bbffaf43f3d637e96"
+checksum = "11181fbabf243db407ef8df94a6ce0b2f9a733bd8be4ad02b4eda9602296cac8"
dependencies = [
- "bitflags 2.6.0",
+ "bitflags 2.9.1",
"errno",
"libc",
- "linux-raw-sys 0.9.3",
- "windows-sys 0.59.0",
+ "linux-raw-sys 0.9.4",
+ "windows-sys 0.60.2",
]
[[package]]
name = "rustls"
-version = "0.23.28"
+version = "0.23.31"
source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "7160e3e10bf4535308537f3c4e1641468cd0e485175d6163087c0393c7d46643"
+checksum = "c0ebcbd2f03de0fc1122ad9bb24b127a5a6cd51d72604a3f3c50ac459762b6cc"
dependencies = [
"aws-lc-rs",
"log",
@@ -3415,9 +3416,9 @@ dependencies = [
[[package]]
name = "rustls-webpki"
-version = "0.103.3"
+version = "0.103.4"
source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "e4a72fe2bcf7a6ac6fd7d0b9e5cb68aeb7d4c0a0271730218b3e92d43b4eb435"
+checksum = "0a17884ae0c1b773f1ccd2bd4a8c72f16da897310a98b0e84bf349ad5ead92fc"
dependencies = [
"aws-lc-rs",
"ring",
@@ -3427,15 +3428,15 @@ dependencies = [
[[package]]
name = "rustversion"
-version = "1.0.18"
+version = "1.0.21"
source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "0e819f2bc632f285be6d7cd36e25940d45b2391dd6d9b939e79de557f7014248"
+checksum = "8a0d197bd2c9dc6e53b84da9556a69ba4cdfab8619eb41a8bd1cc2027a0f6b1d"
[[package]]
name = "ryu"
-version = "1.0.18"
+version = "1.0.20"
source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "f3cb5ba0dc43242ce17de99c180e96db90b235b8a9fdc9543c96d2209116bd9f"
+checksum = "28d3b2b1366ec20994f1fd18c3c594f05c5dd4bc44d8bb0c1c632c8d6829481f"
[[package]]
name = "same-file"
@@ -3454,15 +3455,15 @@ checksum = "94143f37725109f92c262ed2cf5e59bce7498c01bcc1502d7b9afe439a4e9f49"
[[package]]
name = "semver"
-version = "1.0.24"
+version = "1.0.26"
source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "3cb6eb87a131f756572d7fb904f6e7b68633f09cca868c5df1c4b8d1a694bbba"
+checksum = "56e6fa9c48d24d85fb3de5ad847117517440f6beceb7798af16b4a87d616b8d0"
[[package]]
name = "seq-macro"
-version = "0.3.5"
+version = "0.3.6"
source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "a3f0bf26fd526d2a95683cd0f87bf103b8539e2ca1ef48ce002d67aad59aa0b4"
+checksum = "1bc711410fbe7399f390ca1c3b60ad0f53f80e95c5eb935e52268a0e2cd49acc"
[[package]]
name = "serde"
@@ -3481,7 +3482,7 @@ checksum = "5b0276cf7f2c73365f7157c8123c21cd9a50fbbd844757af28ca1f5925fc2a00"
dependencies = [
"proc-macro2",
"quote",
- "syn 2.0.91",
+ "syn 2.0.104",
]
[[package]]
@@ -3519,9 +3520,9 @@ dependencies = [
[[package]]
name = "sha2"
-version = "0.10.8"
+version = "0.10.9"
source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "793db75ad2bcafc3ffa7c68b215fee268f537982cd901d132f89c6343f3a3dc8"
+checksum = "a7507d819769d01a365ab707794a4084392c824f54a7a6a7862f8c3d0892b283"
dependencies = [
"cfg-if",
"cpufeatures",
@@ -3542,24 +3543,21 @@ checksum = "e3a9fe34e3e7a50316060351f37187a3f546bce95496156754b601a5fa71b76e"
[[package]]
name = "siphasher"
-version = "0.3.11"
+version = "1.0.1"
source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "38b58827f4464d87d377d175e90bf58eb00fd8716ff0a62f80356b5e61555d0d"
+checksum = "56199f7ddabf13fe5074ce809e7d3f42b42ae711800501b5b16ea82ad029c39d"
[[package]]
name = "slab"
-version = "0.4.9"
+version = "0.4.10"
source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "8f92a496fb766b417c996b9c5e57daf2f7ad3b0bebe1ccfca4856390e3d3bb67"
-dependencies = [
- "autocfg",
-]
+checksum = "04dc19736151f35336d325007ac991178d504a119863a2fcb3758cdb5e52c50d"
[[package]]
name = "smallvec"
-version = "1.13.2"
+version = "1.15.1"
source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "3c5e1a9a646d36c3599cd173a41282daf47c44583ad367b8e6837255952e5c67"
+checksum = "67b1b7a3b5fe4f1376887184045fcf45c69e92af734b7aaddc05fb777b6fbd03"
[[package]]
name = "smallvec"
@@ -3569,23 +3567,23 @@ checksum = "87b96efa4bd6bdd2ff0c6615cc36fc4970cbae63cfd46ddff5cee35a1b4df570"
[[package]]
name = "snafu"
-version = "0.8.5"
+version = "0.8.6"
source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "223891c85e2a29c3fe8fb900c1fae5e69c2e42415e3177752e8718475efa5019"
+checksum = "320b01e011bf8d5d7a4a4a4be966d9160968935849c83b918827f6a435e7f627"
dependencies = [
"snafu-derive",
]
[[package]]
name = "snafu-derive"
-version = "0.8.5"
+version = "0.8.6"
source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "03c3c6b7927ffe7ecaa769ee0e3994da3b8cafc8f444578982c83ecb161af917"
+checksum = "1961e2ef424c1424204d3a5d6975f934f56b6d50ff5732382d84ebf460e147f7"
dependencies = [
"heck",
"proc-macro2",
"quote",
- "syn 2.0.91",
+ "syn 2.0.104",
]
[[package]]
@@ -3631,14 +3629,14 @@ dependencies = [
"simdutf8",
"sonic-number",
"sonic-simd",
- "thiserror 2.0.9",
+ "thiserror 2.0.12",
]
[[package]]
name = "sonic-simd"
-version = "0.1.0"
+version = "0.1.1"
source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "940a24e82c9a97483ef66cef06b92160a8fa5cd74042c57c10b24d99d169d2fc"
+checksum = "b421f7b6aa4a5de8f685aaf398dfaa828346ee639d2b1c1061ab43d40baa6223"
dependencies = [
"cfg-if",
]
@@ -3670,7 +3668,7 @@ checksum = "01b2e185515564f15375f593fb966b5718bc624ba77fe49fa4616ad619690554"
dependencies = [
"proc-macro2",
"quote",
- "syn 2.0.91",
+ "syn 2.0.104",
]
[[package]]
@@ -3728,7 +3726,7 @@ dependencies = [
"proc-macro2",
"quote",
"rustversion",
- "syn 2.0.91",
+ "syn 2.0.104",
]
[[package]]
@@ -3739,9 +3737,9 @@ checksum = "13c2bddecc57b384dee18652358fb23172facb8a2c51ccc10d74c157bdea3292"
[[package]]
name = "symbolic-common"
-version = "12.13.3"
+version = "12.16.1"
source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "13a4dfe4bbeef59c1f32fc7524ae7c95b9e1de5e79a43ce1604e181081d71b0c"
+checksum = "70f4d06896c59fabe3fe36d7bc003c975f0a0af67d380e14a95eaebffe4f8de5"
dependencies = [
"debugid",
"memmap2",
@@ -3751,9 +3749,9 @@ dependencies = [
[[package]]
name = "symbolic-demangle"
-version = "12.13.3"
+version = "12.16.1"
source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "98cf6a95abff97de4d7ff3473f33cacd38f1ddccad5c1feab435d6760300e3b6"
+checksum = "bd3903bafe2ed4c3512ff4c6eb77cc22b6f43662f3b9f7e3fe4f152927f54ec8"
dependencies = [
"cpp_demangle",
"rustc-demangle",
@@ -3773,9 +3771,9 @@ dependencies = [
[[package]]
name = "syn"
-version = "2.0.91"
+version = "2.0.104"
source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "d53cbcb5a243bd33b7858b1d7f4aca2153490815872d86d955d6ea29f743c035"
+checksum = "17b6f705963418cdb9927482fa304bc562ece2fdd4f616084c50b7023b435a40"
dependencies = [
"proc-macro2",
"quote",
@@ -3793,13 +3791,13 @@ dependencies = [
[[package]]
name = "synstructure"
-version = "0.13.1"
+version = "0.13.2"
source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "c8af7666ab7b6390ab78131fb5b0fce11d6b7a6951602017c35fa82800708971"
+checksum = "728a70f3dbaf5bab7f0c4b1ac8d7ae5ea60a4b5549c8a5914361c99147a709d2"
dependencies = [
"proc-macro2",
"quote",
- "syn 2.0.91",
+ "syn 2.0.104",
]
[[package]]
@@ -3815,9 +3813,9 @@ source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "e8a64e3985349f2441a1a9ef0b853f869006c3855f2cda6862a94d26ebb9d6a1"
dependencies = [
"fastrand",
- "getrandom 0.3.1",
+ "getrandom 0.3.3",
"once_cell",
- "rustix 1.0.3",
+ "rustix 1.0.8",
"windows-sys 0.59.0",
]
@@ -3832,11 +3830,11 @@ dependencies = [
[[package]]
name = "thiserror"
-version = "2.0.9"
+version = "2.0.12"
source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "f072643fd0190df67a8bab670c20ef5d8737177d6ac6b2e9a236cb096206b2cc"
+checksum = "567b8a2dae586314f7be2a752ec7474332959c6460e02bde30d702a66d488708"
dependencies = [
- "thiserror-impl 2.0.9",
+ "thiserror-impl 2.0.12",
]
[[package]]
@@ -3847,18 +3845,18 @@ checksum = "4fee6c4efc90059e10f81e6d42c60a18f76588c3d74cb83a0b242a2b6c7504c1"
dependencies = [
"proc-macro2",
"quote",
- "syn 2.0.91",
+ "syn 2.0.104",
]
[[package]]
name = "thiserror-impl"
-version = "2.0.9"
+version = "2.0.12"
source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "7b50fa271071aae2e6ee85f842e2e28ba8cd2c5fb67f11fcb1fd70b276f9e7d4"
+checksum = "7f7cf42b4507d8ea322120659672cf1b9dbb93f8f2d4ecfd6e51350ff5b17a1d"
dependencies = [
"proc-macro2",
"quote",
- "syn 2.0.91",
+ "syn 2.0.104",
]
[[package]]
@@ -3914,9 +3912,9 @@ dependencies = [
[[package]]
name = "tinystr"
-version = "0.7.6"
+version = "0.8.1"
source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "9117f5d4db391c1cf6927e7bea3db74b9a1c1add8f7eda9ffd5364f40f57b82f"
+checksum = "5d4f6d1145dcb577acf783d4e601bc1d76a13337bb54e6233add580b07344c8b"
dependencies = [
"displaydoc",
"zerovec",
@@ -3924,9 +3922,9 @@ dependencies = [
[[package]]
name = "tinyvec"
-version = "1.8.1"
+version = "1.9.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "022db8904dfa342efe721985167e9fcd16c29b226db4397ed752a761cfce81e8"
+checksum = "09b3661f17e86524eccd4371ab0429194e0d7c008abb45f7a7495b1719463c71"
dependencies = [
"tinyvec_macros",
]
@@ -3963,7 +3961,7 @@ checksum = "6e06d43f1345a3bcd39f6a56dbb7dcab2ba47e68e8ac134855e7e2bdbaf8cab8"
dependencies = [
"proc-macro2",
"quote",
- "syn 2.0.91",
+ "syn 2.0.104",
]
[[package]]
@@ -3989,9 +3987,9 @@ dependencies = [
[[package]]
name = "tokio-util"
-version = "0.7.14"
+version = "0.7.16"
source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "6b9590b93e6fcc1739458317cccd391ad3955e2bde8913edf6f95f9e65a8f034"
+checksum = "14307c986784f72ef81c89db7d9e28d6ac26d16213b109ea501696195e6e3ce5"
dependencies = [
"bytes 1.10.1",
"futures-core",
@@ -4029,7 +4027,7 @@ dependencies = [
"prost-build",
"prost-types",
"quote",
- "syn 2.0.91",
+ "syn 2.0.104",
]
[[package]]
@@ -4045,20 +4043,20 @@ dependencies = [
[[package]]
name = "tracing-attributes"
-version = "0.1.28"
+version = "0.1.30"
source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "395ae124c09f9e6918a2310af6038fba074bcf474ac352496d5910dd59a2226d"
+checksum = "81383ab64e72a7a8b8e13130c49e3dab29def6d0c7d76a03087b3cf71c5c6903"
dependencies = [
"proc-macro2",
"quote",
- "syn 2.0.91",
+ "syn 2.0.104",
]
[[package]]
name = "tracing-core"
-version = "0.1.33"
+version = "0.1.34"
source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "e672c95779cf947c5311f83787af4fa8fffd12fb27e4993211a84bdfd9610f9c"
+checksum = "b9d12581f227e93f094d3af2ae690a574abb8a2b9b7a96e7cfe9647b2b617678"
dependencies = [
"once_cell",
]
@@ -4091,9 +4089,9 @@ checksum = "8b907da542cbced5261bd3256de1b3a1bf340a3d37f93425a07362a1d687de56"
[[package]]
name = "typenum"
-version = "1.17.0"
+version = "1.18.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "42ff0bf0c66b8238c6f3b578df37d0b7848e55df8577b3f74f92a69acceeb825"
+checksum = "1dccffe3ce07af9386bfd29e80c0ab1a8205a2fc34e4bcd40364df902cfa8f3f"
[[package]]
name = "uncased"
@@ -4112,9 +4110,9 @@ checksum = "eeba86d422ce181a719445e51872fa30f1f7413b62becb52e95ec91aa262d85c"
[[package]]
name = "unicode-ident"
-version = "1.0.14"
+version = "1.0.18"
source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "adb9e6ca4f869e1180728b7950e35922a7fc6397f7b641499e8f3ef06e50dc83"
+checksum = "5a5f39404a5da50712a4c1eecf25e90dd62b613502b7e925fd4e4d19b5c96512"
[[package]]
name = "unicode-segmentation"
@@ -4124,9 +4122,9 @@ checksum = "f6ccf251212114b54433ec949fd6a7841275f9ada20dddd2f29e9ceea4501493"
[[package]]
name = "unicode-width"
-version = "0.2.0"
+version = "0.2.1"
source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "1fc81956842c57dac11422a97c3b8195a1ff727f06e85c84ed2e8aa277c9a0fd"
+checksum = "4a1a07cc7db3810833284e8d372ccdc6da29741639ecc70c9ec107df0fa6154c"
[[package]]
name = "untrusted"
@@ -4145,12 +4143,6 @@ dependencies = [
"percent-encoding",
]
-[[package]]
-name = "utf16_iter"
-version = "1.0.5"
-source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "c8232dd3cdaed5356e0f716d285e4b40b932ac434100fe9b7e0e8e935b9e6246"
-
[[package]]
name = "utf8_iter"
version = "1.0.4"
@@ -4163,7 +4155,7 @@ version = "1.17.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "3cf4199d1e5d15ddd86a694e4d0dffa9c323ce759fea589f00fef9d81cc1931d"
dependencies = [
- "getrandom 0.3.1",
+ "getrandom 0.3.3",
"js-sys",
"wasm-bindgen",
]
@@ -4186,15 +4178,15 @@ dependencies = [
[[package]]
name = "wasi"
-version = "0.11.0+wasi-snapshot-preview1"
+version = "0.11.1+wasi-snapshot-preview1"
source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "9c8d87e72b64a3b4db28d11ce29237c246188f4f51057d65a7eab63b7987e423"
+checksum = "ccf3ec651a847eb01de73ccad15eb7d99f80485de043efb2f370cd654f4ea44b"
[[package]]
name = "wasi"
-version = "0.13.3+wasi-0.2.2"
+version = "0.14.2+wasi-0.2.4"
source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "26816d2e1a4a36a2940b96c5296ce403917633dff8f3440e9b236ed6f6bacad2"
+checksum = "9683f9a5a998d873c0d21fcbe3c083009670149a8fab228644b8bd36b2c48cb3"
dependencies = [
"wit-bindgen-rt",
]
@@ -4221,7 +4213,7 @@ dependencies = [
"log",
"proc-macro2",
"quote",
- "syn 2.0.91",
+ "syn 2.0.104",
"wasm-bindgen-shared",
]
@@ -4243,7 +4235,7 @@ checksum = "8ae87ea40c9f689fc23f209965b6fb8a99ad69aeeb0231408be24920604395de"
dependencies = [
"proc-macro2",
"quote",
- "syn 2.0.91",
+ "syn 2.0.104",
"wasm-bindgen-backend",
"wasm-bindgen-shared",
]
@@ -4318,11 +4310,61 @@ checksum = "712e227841d057c1ee1cd2fb22fa7e5a5461ae8e48fa2ca79ec42cfc1931183f"
[[package]]
name = "windows-core"
-version = "0.52.0"
+version = "0.61.2"
source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "33ab640c8d7e35bf8ba19b884ba838ceb4fba93a4e8c65a9059d08afcfc683d9"
+checksum = "c0fdd3ddb90610c7638aa2b3a3ab2904fb9e5cdbecc643ddb3647212781c4ae3"
dependencies = [
- "windows-targets",
+ "windows-implement",
+ "windows-interface",
+ "windows-link",
+ "windows-result",
+ "windows-strings",
+]
+
+[[package]]
+name = "windows-implement"
+version = "0.60.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "a47fddd13af08290e67f4acabf4b459f647552718f683a7b415d290ac744a836"
+dependencies = [
+ "proc-macro2",
+ "quote",
+ "syn 2.0.104",
+]
+
+[[package]]
+name = "windows-interface"
+version = "0.59.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "bd9211b69f8dcdfa817bfd14bf1c97c9188afa36f4750130fcdf3f400eca9fa8"
+dependencies = [
+ "proc-macro2",
+ "quote",
+ "syn 2.0.104",
+]
+
+[[package]]
+name = "windows-link"
+version = "0.1.3"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "5e6ad25900d524eaabdbbb96d20b4311e1e7ae1699af4fb28c17ae66c80d798a"
+
+[[package]]
+name = "windows-result"
+version = "0.3.4"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "56f42bd332cc6c8eac5af113fc0c1fd6a8fd2aa08a0119358686e5160d0586c6"
+dependencies = [
+ "windows-link",
+]
+
+[[package]]
+name = "windows-strings"
+version = "0.4.2"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "56e6c93f3a0c3b36176cb1327a4958a0353d5d166c2a35cb268ace15e91d3b57"
+dependencies = [
+ "windows-link",
]
[[package]]
@@ -4331,7 +4373,7 @@ version = "0.52.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "282be5f36a8ce781fad8c8ae18fa3f9beff57ec1b52cb3de0789201425d9a33d"
dependencies = [
- "windows-targets",
+ "windows-targets 0.52.6",
]
[[package]]
@@ -4340,7 +4382,16 @@ version = "0.59.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "1e38bc4d79ed67fd075bcc251a1c39b32a1776bbe92e5bef1f0bf1f8c531853b"
dependencies = [
- "windows-targets",
+ "windows-targets 0.52.6",
+]
+
+[[package]]
+name = "windows-sys"
+version = "0.60.2"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "f2f500e4d28234f72040990ec9d39e3a6b950f9f22d3dba18416c35882612bcb"
+dependencies = [
+ "windows-targets 0.53.3",
]
[[package]]
@@ -4349,14 +4400,31 @@ version = "0.52.6"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "9b724f72796e036ab90c1021d4780d4d3d648aca59e491e6b98e725b84e99973"
dependencies = [
- "windows_aarch64_gnullvm",
- "windows_aarch64_msvc",
- "windows_i686_gnu",
- "windows_i686_gnullvm",
- "windows_i686_msvc",
- "windows_x86_64_gnu",
- "windows_x86_64_gnullvm",
- "windows_x86_64_msvc",
+ "windows_aarch64_gnullvm 0.52.6",
+ "windows_aarch64_msvc 0.52.6",
+ "windows_i686_gnu 0.52.6",
+ "windows_i686_gnullvm 0.52.6",
+ "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",
+]
+
+[[package]]
+name = "windows-targets"
+version = "0.53.3"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "d5fe6031c4041849d7c496a8ded650796e7b6ecc19df1a431c1a363342e5dc91"
+dependencies = [
+ "windows-link",
+ "windows_aarch64_gnullvm 0.53.0",
+ "windows_aarch64_msvc 0.53.0",
+ "windows_i686_gnu 0.53.0",
+ "windows_i686_gnullvm 0.53.0",
+ "windows_i686_msvc 0.53.0",
+ "windows_x86_64_gnu 0.53.0",
+ "windows_x86_64_gnullvm 0.53.0",
+ "windows_x86_64_msvc 0.53.0",
]
[[package]]
@@ -4365,77 +4433,119 @@ version = "0.52.6"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "32a4622180e7a0ec044bb555404c800bc9fd9ec262ec147edd5989ccd0c02cd3"
+[[package]]
+name = "windows_aarch64_gnullvm"
+version = "0.53.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "86b8d5f90ddd19cb4a147a5fa63ca848db3df085e25fee3cc10b39b6eebae764"
+
[[package]]
name = "windows_aarch64_msvc"
version = "0.52.6"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "09ec2a7bb152e2252b53fa7803150007879548bc709c039df7627cabbd05d469"
+[[package]]
+name = "windows_aarch64_msvc"
+version = "0.53.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "c7651a1f62a11b8cbd5e0d42526e55f2c99886c77e007179efff86c2b137e66c"
+
[[package]]
name = "windows_i686_gnu"
version = "0.52.6"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "8e9b5ad5ab802e97eb8e295ac6720e509ee4c243f69d781394014ebfe8bbfa0b"
+[[package]]
+name = "windows_i686_gnu"
+version = "0.53.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "c1dc67659d35f387f5f6c479dc4e28f1d4bb90ddd1a5d3da2e5d97b42d6272c3"
+
[[package]]
name = "windows_i686_gnullvm"
version = "0.52.6"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "0eee52d38c090b3caa76c563b86c3a4bd71ef1a819287c19d586d7334ae8ed66"
+[[package]]
+name = "windows_i686_gnullvm"
+version = "0.53.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "9ce6ccbdedbf6d6354471319e781c0dfef054c81fbc7cf83f338a4296c0cae11"
+
[[package]]
name = "windows_i686_msvc"
version = "0.52.6"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "240948bc05c5e7c6dabba28bf89d89ffce3e303022809e73deaefe4f6ec56c66"
+[[package]]
+name = "windows_i686_msvc"
+version = "0.53.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "581fee95406bb13382d2f65cd4a908ca7b1e4c2f1917f143ba16efe98a589b5d"
+
[[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_gnu"
+version = "0.53.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "2e55b5ac9ea33f2fc1716d1742db15574fd6fc8dadc51caab1c16a3d3b4190ba"
+
[[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_gnullvm"
+version = "0.53.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "0a6e035dd0599267ce1ee132e51c27dd29437f63325753051e71dd9e42406c57"
+
[[package]]
name = "windows_x86_64_msvc"
version = "0.52.6"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "589f6da84c646204747d1270a2a5661ea66ed1cced2631d546fdfb155959f9ec"
+[[package]]
+name = "windows_x86_64_msvc"
+version = "0.53.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "271414315aff87387382ec3d271b52d7ae78726f5d44ac98b4f4030c91880486"
+
[[package]]
name = "winnow"
-version = "0.7.11"
+version = "0.7.12"
source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "74c7b26e3480b707944fc872477815d29a8e429d2f93a1ce000f5fa84a15cbcd"
+checksum = "f3edebf492c8125044983378ecb5766203ad3b4c2f7a922bd7dd207f6d443e95"
dependencies = [
"memchr",
]
[[package]]
name = "wit-bindgen-rt"
-version = "0.33.0"
+version = "0.39.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "3268f3d866458b787f390cf61f4bbb563b922d091359f9608842999eaee3943c"
+checksum = "6f42320e61fe2cfd34354ecb597f86f413484a798ba44a8ca1165c58d42da6c1"
dependencies = [
- "bitflags 2.6.0",
+ "bitflags 2.9.1",
]
-[[package]]
-name = "write16"
-version = "1.0.0"
-source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "d1890f4022759daae28ed4fe62859b1236caebfc61ede2f63ed4e695f3f6d936"
-
[[package]]
name = "writeable"
-version = "0.5.5"
+version = "0.6.1"
source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "1e9df38ee2d2c3c5948ea468a8406ff0db0b29ae1ffde1bcf20ef305bcc95c51"
+checksum = "ea2f10b9bb0928dfb1b42b65e1f9e36f7f54dbdf08457afefb38afcdec4fa2bb"
[[package]]
name = "wyz"
@@ -4457,9 +4567,9 @@ dependencies = [
[[package]]
name = "yoke"
-version = "0.7.5"
+version = "0.8.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "120e6aef9aa629e3d4f52dc8cc43a015c7724194c97dfaf45180d2daf2b77f40"
+checksum = "5f41bb01b8226ef4bfd589436a297c53d118f65921786300e427be8d487695cc"
dependencies = [
"serde",
"stable_deref_trait",
@@ -4469,75 +4579,54 @@ dependencies = [
[[package]]
name = "yoke-derive"
-version = "0.7.5"
+version = "0.8.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "2380878cad4ac9aac1e2435f3eb4020e8374b5f13c296cb75b4620ff8e229154"
+checksum = "38da3c9736e16c5d3c8c597a9aaa5d1fa565d0532ae05e27c24aa62fb32c0ab6"
dependencies = [
"proc-macro2",
"quote",
- "syn 2.0.91",
+ "syn 2.0.104",
"synstructure",
]
[[package]]
name = "zerocopy"
-version = "0.7.35"
-source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "1b9b4fd18abc82b8136838da5d50bae7bdea537c574d8dc1a34ed098d6c166f0"
-dependencies = [
- "byteorder",
- "zerocopy-derive 0.7.35",
-]
-
-[[package]]
-name = "zerocopy"
-version = "0.8.14"
+version = "0.8.26"
source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "a367f292d93d4eab890745e75a778da40909cab4d6ff8173693812f79c4a2468"
+checksum = "1039dd0d3c310cf05de012d8a39ff557cb0d23087fd44cad61df08fc31907a2f"
dependencies = [
- "zerocopy-derive 0.8.14",
+ "zerocopy-derive",
]
[[package]]
name = "zerocopy-derive"
-version = "0.7.35"
+version = "0.8.26"
source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "fa4f8080344d4671fb4e831a13ad1e68092748387dfc4f55e356242fae12ce3e"
+checksum = "9ecf5b4cc5364572d7f4c329661bcc82724222973f2cab6f050a4e5c22f75181"
dependencies = [
"proc-macro2",
"quote",
- "syn 2.0.91",
-]
-
-[[package]]
-name = "zerocopy-derive"
-version = "0.8.14"
-source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "d3931cb58c62c13adec22e38686b559c86a30565e16ad6e8510a337cedc611e1"
-dependencies = [
- "proc-macro2",
- "quote",
- "syn 2.0.91",
+ "syn 2.0.104",
]
[[package]]
name = "zerofrom"
-version = "0.1.5"
+version = "0.1.6"
source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "cff3ee08c995dee1859d998dea82f7374f2826091dd9cd47def953cae446cd2e"
+checksum = "50cc42e0333e05660c3587f3bf9d0478688e15d870fab3346451ce7f8c9fbea5"
dependencies = [
"zerofrom-derive",
]
[[package]]
name = "zerofrom-derive"
-version = "0.1.5"
+version = "0.1.6"
source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "595eed982f7d355beb85837f651fa22e90b3c044842dc7f2c2842c086f295808"
+checksum = "d71e5d6e06ab090c67b5e44993ec16b72dcbaabc526db883a360057678b48502"
dependencies = [
"proc-macro2",
"quote",
- "syn 2.0.91",
+ "syn 2.0.104",
"synstructure",
]
@@ -4547,11 +4636,22 @@ version = "1.8.1"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "ced3678a2879b30306d323f4542626697a464a97c0a07c9aebf7ebca65cd4dde"
+[[package]]
+name = "zerotrie"
+version = "0.2.2"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "36f0bbd478583f79edad978b407914f61b2972f5af6fa089686016be8f9af595"
+dependencies = [
+ "displaydoc",
+ "yoke",
+ "zerofrom",
+]
+
[[package]]
name = "zerovec"
-version = "0.10.4"
+version = "0.11.3"
source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "aa2b893d79df23bfb12d5461018d408ea19dfafe76c2c7ef6d4eba614f8ff079"
+checksum = "bdbb9122ea75b11bf96e7492afb723e8a7fbe12c67417aa95e7e3d18144d37cd"
dependencies = [
"yoke",
"zerofrom",
@@ -4560,13 +4660,13 @@ dependencies = [
[[package]]
name = "zerovec-derive"
-version = "0.10.3"
+version = "0.11.1"
source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "6eafa6dfb17584ea3e2bd6e76e0cc15ad7af12b09abdd1ca55961bed9b1063c6"
+checksum = "5b96237efa0c878c64bd89c436f661be4e46b2f3eff1ebb976f7ef2321d2f58f"
dependencies = [
"proc-macro2",
"quote",
- "syn 2.0.91",
+ "syn 2.0.104",
]
[[package]]
diff --git a/Cargo.toml b/Cargo.toml
index 6d0b03be..d80249d2 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -72,7 +72,7 @@ bitvec = "1.0.1"
byteorder = "1.5.0"
bytes = "1.10.1"
bytesize = "2.0.1"
-chrono = "0.4.33"
+chrono = "=0.4.39"
count-write = "0.1.0"
foldhash = "0.1.5"
futures = "0.3"
@@ -102,27 +102,27 @@ zstd = "0.13.3"
[patch.crates-io]
# datafusion: branch=v42-blaze
-datafusion = { git = "https://github.com/blaze-init/arrow-datafusion.git", rev = "b7876646a02401cc064657ccf17de27cfe6240ee"}
-datafusion-common = { git = "https://github.com/blaze-init/arrow-datafusion.git", rev = "b7876646a02401cc064657ccf17de27cfe6240ee"}
-datafusion-expr = { git = "https://github.com/blaze-init/arrow-datafusion.git", rev = "b7876646a02401cc064657ccf17de27cfe6240ee"}
-datafusion-execution = { git = "https://github.com/blaze-init/arrow-datafusion.git", rev = "b7876646a02401cc064657ccf17de27cfe6240ee"}
-datafusion-optimizer = { git = "https://github.com/blaze-init/arrow-datafusion.git", rev = "b7876646a02401cc064657ccf17de27cfe6240ee"}
-datafusion-physical-expr = { git = "https://github.com/blaze-init/arrow-datafusion.git", rev = "b7876646a02401cc064657ccf17de27cfe6240ee"}
-orc-rust = { git = "https://github.com/blaze-init/datafusion-orc.git", rev = "172183facb1d05d6ac0b4ec234207a52387db2a3"}
+datafusion = { git = "https://github.com/blaze-init/arrow-datafusion.git", rev = "0e6a0b0942cd052d8c2a3606f26d9823afc911d9"}
+datafusion-common = { git = "https://github.com/blaze-init/arrow-datafusion.git", rev = "0e6a0b0942cd052d8c2a3606f26d9823afc911d9"}
+datafusion-expr = { git = "https://github.com/blaze-init/arrow-datafusion.git", rev = "0e6a0b0942cd052d8c2a3606f26d9823afc911d9"}
+datafusion-execution = { git = "https://github.com/blaze-init/arrow-datafusion.git", rev = "0e6a0b0942cd052d8c2a3606f26d9823afc911d9"}
+datafusion-optimizer = { git = "https://github.com/blaze-init/arrow-datafusion.git", rev = "0e6a0b0942cd052d8c2a3606f26d9823afc911d9"}
+datafusion-physical-expr = { git = "https://github.com/blaze-init/arrow-datafusion.git", rev = "0e6a0b0942cd052d8c2a3606f26d9823afc911d9"}
+orc-rust = { git = "https://github.com/blaze-init/datafusion-orc.git", rev = "6499dffa74267623ee311b1ca68db34237e83b00"}
# arrow: branch=v53-blaze
-arrow = { git = "https://github.com/blaze-init/arrow-rs.git", rev = "1fbdc9f980aa34c7758da9db79a64248f24daf1e"}
-arrow-arith = { git = "https://github.com/blaze-init/arrow-rs.git", rev = "1fbdc9f980aa34c7758da9db79a64248f24daf1e"}
-arrow-array = { git = "https://github.com/blaze-init/arrow-rs.git", rev = "1fbdc9f980aa34c7758da9db79a64248f24daf1e"}
-arrow-buffer = { git = "https://github.com/blaze-init/arrow-rs.git", rev = "1fbdc9f980aa34c7758da9db79a64248f24daf1e"}
-arrow-cast = { git = "https://github.com/blaze-init/arrow-rs.git", rev = "1fbdc9f980aa34c7758da9db79a64248f24daf1e"}
-arrow-data = { git = "https://github.com/blaze-init/arrow-rs.git", rev = "1fbdc9f980aa34c7758da9db79a64248f24daf1e"}
-arrow-ord = { git = "https://github.com/blaze-init/arrow-rs.git", rev = "1fbdc9f980aa34c7758da9db79a64248f24daf1e"}
-arrow-row = { git = "https://github.com/blaze-init/arrow-rs.git", rev = "1fbdc9f980aa34c7758da9db79a64248f24daf1e"}
-arrow-schema = { git = "https://github.com/blaze-init/arrow-rs.git", rev = "1fbdc9f980aa34c7758da9db79a64248f24daf1e"}
-arrow-select = { git = "https://github.com/blaze-init/arrow-rs.git", rev = "1fbdc9f980aa34c7758da9db79a64248f24daf1e"}
-arrow-string = { git = "https://github.com/blaze-init/arrow-rs.git", rev = "1fbdc9f980aa34c7758da9db79a64248f24daf1e"}
-parquet = { git = "https://github.com/blaze-init/arrow-rs.git", rev = "1fbdc9f980aa34c7758da9db79a64248f24daf1e"}
+arrow = { git = "https://github.com/blaze-init/arrow-rs.git", rev = "94ef6f608254484c0ecfc33fcef19891c4f87612"}
+arrow-arith = { git = "https://github.com/blaze-init/arrow-rs.git", rev = "94ef6f608254484c0ecfc33fcef19891c4f87612"}
+arrow-array = { git = "https://github.com/blaze-init/arrow-rs.git", rev = "94ef6f608254484c0ecfc33fcef19891c4f87612"}
+arrow-buffer = { git = "https://github.com/blaze-init/arrow-rs.git", rev = "94ef6f608254484c0ecfc33fcef19891c4f87612"}
+arrow-cast = { git = "https://github.com/blaze-init/arrow-rs.git", rev = "94ef6f608254484c0ecfc33fcef19891c4f87612"}
+arrow-data = { git = "https://github.com/blaze-init/arrow-rs.git", rev = "94ef6f608254484c0ecfc33fcef19891c4f87612"}
+arrow-ord = { git = "https://github.com/blaze-init/arrow-rs.git", rev = "94ef6f608254484c0ecfc33fcef19891c4f87612"}
+arrow-row = { git = "https://github.com/blaze-init/arrow-rs.git", rev = "94ef6f608254484c0ecfc33fcef19891c4f87612"}
+arrow-schema = { git = "https://github.com/blaze-init/arrow-rs.git", rev = "94ef6f608254484c0ecfc33fcef19891c4f87612"}
+arrow-select = { git = "https://github.com/blaze-init/arrow-rs.git", rev = "94ef6f608254484c0ecfc33fcef19891c4f87612"}
+arrow-string = { git = "https://github.com/blaze-init/arrow-rs.git", rev = "94ef6f608254484c0ecfc33fcef19891c4f87612"}
+parquet = { git = "https://github.com/blaze-init/arrow-rs.git", rev = "94ef6f608254484c0ecfc33fcef19891c4f87612"}
# serde_json: branch=v1.0.96-blaze
serde_json = { git = "https://github.com/blaze-init/json", branch = "v1.0.96-blaze" }
diff --git a/build-native.sh b/build-native.sh
index 67d68d9b..30e29c80 100755
--- a/build-native.sh
+++ b/build-native.sh
@@ -50,7 +50,7 @@ echo -e "old build-checksum: \n$old_checksum\n========"
echo -e "new build-checksum: \n$new_checksum\n========"
if [ "$new_checksum" != "$old_checksum" ]; then
- export RUSTFLAGS=${RUSTFLAGS:-"-C target-cpu=native"}
+ export RUSTFLAGS=${RUSTFLAGS:-"-C target-cpu=skylake"}
echo "Running cargo fix..."
cargo fix --all --allow-dirty --allow-staged --allow-no-vcs
diff --git a/dev/mvn-build-helper/assembly/pom.xml b/dev/mvn-build-helper/assembly/pom.xml
index dae047e9..b0f3575e 100644
--- a/dev/mvn-build-helper/assembly/pom.xml
+++ b/dev/mvn-build-helper/assembly/pom.xml
@@ -110,7 +110,7 @@
<archive>
<addMavenDescriptor>false</addMavenDescriptor>
</archive>
- <outputFile>../../../target/blaze-engine-${shimName}-${releaseMode}-${revision}.jar</outputFile>
+ <outputFile>../../../target/blaze-engine-${shimName}-${releaseMode}-${revision}${pkgSuffix}.jar</outputFile>
</configuration>
</execution>
</executions>
diff --git a/pom.xml b/pom.xml
index 62ec7572..b455cdac 100644
--- a/pom.xml
+++ b/pom.xml
@@ -16,7 +16,7 @@
</modules>
<properties>
- <revision>5.0.0-SNAPSHOT</revision>
+ <revision>6.0.0-SNAPSHOT</revision>
<project.build.sourceEncoding>UTF-8</project.build.sourceEncoding>
<arrowVersion>16.0.0</arrowVersion>
<protobufVersion>3.21.9</protobufVersion>
@@ -269,6 +269,7 @@
<version>0.1.4_0.9.31</version>
<configuration>
<config>${session.executionRootDirectory}/scalafix.conf</config>
+ <skipTest>true</skipTest>
</configuration>
<dependencies>
<dependency>
@@ -318,10 +319,28 @@
</properties>
</profile>
+ <profile>
+ <id>spark-241kwaiae</id>
+ <properties>
+ <shimName>spark-241kwaiae</shimName>
+ <pkgSuffix />
+ <shimPkg>spark-extension-shims-spark241kwaiae</shimPkg>
+ <javaVersion>1.8</javaVersion>
+ <scalaVersion>2.11</scalaVersion>
+ <scalaLongVersion>2.11.12</scalaLongVersion>
+ <scalaTestVersion>3.0.8</scalaTestVersion>
+ <scalafmtVersion>3.0.0-RC1</scalafmtVersion>
+ <celebornVersion>0.5.2</celebornVersion>
+ <sparkVersion>2.4.1U1.9.44-master-BETA-SNAPSHOT</sparkVersion>
+ <kwaiRssShuffleSdkVersion>1.3.75-RELEASE</kwaiRssShuffleSdkVersion>
+ </properties>
+ </profile>
+
<profile>
<id>spark-3.0</id>
<properties>
<shimName>spark-3.0</shimName>
+ <pkgSuffix />
<shimPkg>spark-extension-shims-spark3</shimPkg>
<javaVersion>1.8</javaVersion>
<scalaVersion>2.12</scalaVersion>
@@ -336,6 +355,7 @@
<id>spark-3.1</id>
<properties>
<shimName>spark-3.1</shimName>
+ <pkgSuffix />
<shimPkg>spark-extension-shims-spark3</shimPkg>
<javaVersion>1.8</javaVersion>
<scalaVersion>2.12</scalaVersion>
@@ -350,6 +370,7 @@
<id>spark-3.2</id>
<properties>
<shimName>spark-3.2</shimName>
+ <pkgSuffix />
<shimPkg>spark-extension-shims-spark3</shimPkg>
<javaVersion>1.8</javaVersion>
<scalaVersion>2.12</scalaVersion>
@@ -364,6 +385,7 @@
<id>spark-3.3</id>
<properties>
<shimName>spark-3.3</shimName>
+ <pkgSuffix />
<shimPkg>spark-extension-shims-spark3</shimPkg>
<javaVersion>1.8</javaVersion>
<scalaVersion>2.12</scalaVersion>
@@ -378,6 +400,7 @@
<id>spark-3.4</id>
<properties>
<shimName>spark-3.4</shimName>
+ <pkgSuffix />
<shimPkg>spark-extension-shims-spark3</shimPkg>
<javaVersion>1.8</javaVersion>
<scalaVersion>2.12</scalaVersion>
@@ -388,17 +411,19 @@
</properties>
</profile>
+ <!-- modified for kwai package -->
<profile>
<id>spark-3.5</id>
<properties>
<shimName>spark-3.5</shimName>
+ <pkgSuffix>-kwai</pkgSuffix>
<shimPkg>spark-extension-shims-spark3</shimPkg>
<javaVersion>1.8</javaVersion>
<scalaVersion>2.12</scalaVersion>
<scalaLongVersion>2.12.15</scalaLongVersion>
<scalaTestVersion>3.2.9</scalaTestVersion>
<scalafmtVersion>3.0.0</scalafmtVersion>
- <sparkVersion>3.5.5</sparkVersion>
+ <sparkVersion>3.5.1U0.5-SNAPSHOT</sparkVersion>
</properties>
</profile>
</profiles>
diff --git a/spark-extension-shims-spark241kwaiae/pom.xml b/spark-extension-shims-spark241kwaiae/pom.xml
new file mode 100644
index 00000000..622658ee
--- /dev/null
+++ b/spark-extension-shims-spark241kwaiae/pom.xml
@@ -0,0 +1,66 @@
+<?xml version="1.0" encoding="UTF-8"?>
+<project xmlns="http://maven.apache.org/POM/4.0.0" xmlns:xsi="http://www.w3.org/2001/XMLSchema-instance" xsi:schemaLocation="http://maven.apache.org/POM/4.0.0 http://maven.apache.org/xsd/maven-4.0.0.xsd">
+ <modelVersion>4.0.0</modelVersion>
+
+ <parent>
+ <groupId>org.blaze</groupId>
+ <artifactId>blaze-engine</artifactId>
+ <version>${revision}</version>
+ <relativePath>../</relativePath>
+ </parent>
+ <groupId>org.blaze</groupId>
+ <artifactId>spark-extension-shims-spark241kwaiae</artifactId>
+ <packaging>jar</packaging>
+
+ <dependencies>
+ <dependency>
+ <groupId>org.blaze</groupId>
+ <artifactId>spark-extension</artifactId>
+ <version>${revision}</version>
+ </dependency>
+ <dependency>
+ <groupId>org.scala-lang</groupId>
+ <artifactId>scala-library</artifactId>
+ <scope>provided</scope>
+ </dependency>
+ <dependency>
+ <groupId>org.apache.spark</groupId>
+ <artifactId>spark-core_${scalaVersion}</artifactId>
+ <scope>provided</scope>
+ </dependency>
+ <dependency>
+ <groupId>org.apache.spark</groupId>
+ <artifactId>spark-hive_${scalaVersion}</artifactId>
+ <scope>provided</scope>
+ </dependency>
+ <dependency>
+ <groupId>org.apache.spark</groupId>
+ <artifactId>spark-sql_${scalaVersion}</artifactId>
+ <scope>provided</scope>
+ </dependency>
+ <dependency>
+ <groupId>com.kuaishou.dataarch</groupId>
+ <artifactId>shuffle-sdk</artifactId>
+ <version>${kwaiRssShuffleSdkVersion}</version>
+ <classifier>shaded</classifier>
+ <scope>provided</scope>
+ </dependency>
+ <dependency>
+ <groupId>org.apache.arrow</groupId>
+ <artifactId>arrow-c-data</artifactId>
+ </dependency>
+ <dependency>
+ <groupId>org.apache.arrow</groupId>
+ <artifactId>arrow-compression</artifactId>
+ </dependency>
+ <dependency>
+ <groupId>org.apache.arrow</groupId>
+ <artifactId>arrow-vector</artifactId>
+ </dependency>
+ <dependency>
+ <groupId>org.scalatest</groupId>
+ <artifactId>scalatest_${scalaVersion}</artifactId>
+ <scope>test</scope>
+ </dependency>
+ </dependencies>
+</project>
diff --git a/spark-extension-shims-spark241kwaiae/src/main/scala/org/apache/spark/sql/blaze/ShimsImpl.scala b/spark-extension-shims-spark241kwaiae/src/main/scala/org/apache/spark/sql/blaze/ShimsImpl.scala
new file mode 100644
index 00000000..60b203e6
--- /dev/null
+++ b/spark-extension-shims-spark241kwaiae/src/main/scala/org/apache/spark/sql/blaze/ShimsImpl.scala
@@ -0,0 +1,813 @@
+/*
+ * Copyright 2022 The Blaze Authors
+ *
+ * Licensed 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.blaze
+
+import java.io.File
+import java.util.UUID
+import com.kuaishou.dataarch.shuffle.proto.dto.common.PartitionStatistics
+import org.apache.commons.lang3.reflect.FieldUtils
+import org.apache.commons.lang3.reflect.MethodUtils
+import org.apache.spark.{OneToOneDependency, Partition, ShuffleDependency, SparkConf, SparkContext, SparkEnv, SparkException, TaskContext}
+import org.blaze.{protobuf => pb}
+import org.apache.spark.internal.{Logging, config}
+import org.apache.spark.internal.config.OptionalConfigEntry
+import org.apache.spark.rdd.RDD
+import org.apache.spark.rdd.ShuffledRDDPartition
+import org.apache.spark.scheduler.MapStatus
+import org.apache.spark.shuffle.BaseShuffleHandle
+import org.apache.spark.shuffle.IndexShuffleBlockResolver
+import org.apache.spark.shuffle.ShuffleHandle
+import org.apache.spark.shuffle.ShuffleReader
+import org.apache.spark.shuffle.ShuffleWriteMetricsReporter
+import org.apache.spark.shuffle.UnifiedShuffleManager
+import org.apache.spark.shuffle.sort.MapInfo
+import org.apache.spark.sql.catalyst.expressions.Expression
+import org.apache.spark.sql.catalyst.expressions.aggregate.AggregateExpression
+import org.apache.spark.sql.catalyst.expressions.aggregate.First
+import org.apache.spark.sql.catalyst.expressions.Literal
+import org.apache.spark.sql.catalyst.expressions.StringSplit
+import org.apache.spark.sql.catalyst.plans.physical.BroadcastMode
+import org.apache.spark.sql.catalyst.plans.physical.Partitioning
+import org.apache.spark.sql.execution.SparkPlan
+import org.apache.spark.sql.execution.adaptive.BroadcastQueryStageInput
+import org.apache.spark.sql.execution.adaptive.LocalShuffledRowRDD
+import org.apache.spark.sql.execution.adaptive.QueryStage
+import org.apache.spark.sql.execution.adaptive.QueryStageInput
+import org.apache.spark.sql.execution.adaptive.ShuffleQueryStageInput
+import org.apache.spark.sql.execution.adaptive.SkewedShuffleQueryStageInput
+import org.apache.spark.sql.execution.blaze.plan.{BroadcastSide, BuildSide, ConvertToNativeBase, ConvertToNativeExec, NativeAggBase, NativeAggExec, NativeBroadcastExchangeBase, NativeBroadcastExchangeExec, NativeBroadcastJoinBase, NativeExpandBase, NativeExpandExec, NativeFilterBase, NativeFilterExec, NativeGenerateBase, NativeGenerateExec, NativeGlobalLimitBase, NativeGlobalLimitExec, NativeLocalLimitBase, NativeLocalLimitExec, NativeOrcScanBase, NativeOrcScanExec, NativeParquetInsertIn [...]
+import org.apache.spark.sql.execution.blaze.shuffle.BlazeBlockStoreShuffleReaderBase
+import org.apache.spark.sql.execution.exchange.ReusedExchangeExec
+import org.apache.spark.sql.execution.metric.SQLShuffleReadMetricsReporter
+import org.apache.spark.sql.execution.FileSourceScanExec
+import org.apache.spark.sql.execution.blaze.shuffle.RssPartitionWriter
+import org.apache.spark.sql.execution.exchange.BroadcastExchangeLike
+import org.apache.spark.sql.execution.UnaryExecNode
+import org.apache.spark.sql.types.StringType
+import org.apache.spark.sql.SQLContext
+import org.apache.spark.sql.blaze.BlazeConverters.ForceNativeExecutionWrapperBase
+import org.apache.spark.sql.blaze.NativeConverters.NativeExprWrapperBase
+import org.apache.spark.sql.catalyst.expressions.Attribute
+import org.apache.spark.sql.catalyst.expressions.Generator
+import org.apache.spark.sql.catalyst.expressions.NamedExpression
+import org.apache.spark.sql.catalyst.expressions.SortOrder
+import org.apache.spark.sql.catalyst.plans.JoinType
+import org.apache.spark.sql.execution.blaze.plan.NativeAggBase.AggExecMode
+import org.apache.spark.sql.execution.metric.SQLMetric
+import org.apache.spark.sql.hive.execution.InsertIntoHiveTable
+import org.apache.spark.sql.types.DataType
+import org.apache.spark.sql.SparkSession
+import org.apache.spark.sql.SparkSessionExtensions
+import org.apache.spark.sql.catalyst.catalog.CatalogTable
+import org.apache.spark.sql.catalyst.catalog.HiveTableRelation
+import org.apache.spark.sql.catalyst.expressions.PromotePrecision
+import org.apache.spark.sql.catalyst.plans.logical.LogicalPlan
+import org.apache.spark.sql.catalyst.rules.Rule
+import org.apache.spark.sql.execution.datasources.HadoopFsRelation
+import org.apache.spark.sql.execution.datasources.LogicalRelation
+import org.apache.spark.sql.execution.datasources.orc.OrcFileFormat
+import org.apache.spark.sql.execution.datasources.parquet.ParquetFileFormat
+import org.apache.spark.sql.execution.joins.blaze.plan.NativeBroadcastJoinExec
+import org.apache.spark.sql.execution.joins.blaze.plan.NativeShuffledHashJoinExec
+import org.apache.spark.sql.execution.ui.{BlazePlanListener, SQLAppStatusListener, SQLAppStatusStore}
+import org.apache.spark.sql.internal.SQLConf
+import org.apache.spark.status.ElementTrackingStore
+import org.apache.spark.storage.BlockManagerId
+import org.apache.spark.storage.FileSegment
+import org.apache.spark.storage.ShuffleDataBlockId
+import org.apache.spark.util.CompletionIterator
+import org.apache.spark.util.ExternalBlockStoreUtils
+import org.apache.spark.util.Utils
+import org.blaze.protobuf.PhysicalExprNode
+
+class ShimsImpl extends Shims with Logging {
+ override def shimVersion: String = "241kwaiae"
+
+ override def initExtension(): Unit = {
+ val conf = SparkEnv.get.conf
+ val executorMemoryLimitEnable =
+ conf.getBoolean("spark.blaze.executor.memory.limit.enable", false)
+ val executorMemoryMiB = if (!executorMemoryLimitEnable) {
+ conf.get(config.EXECUTOR_MEMORY)
+ } else {
+ conf.getLong("spark.blaze.executor.memory.limit.mb", conf.get(config.EXECUTOR_MEMORY))
+ }
+ conf.set(config.EXECUTOR_MEMORY, executorMemoryMiB)
+ }
+
+ override def onApplyingExtension(extension: SparkSessionExtensions): Unit = {
+ extension.injectOptimizerRule(sparkSession => {
+ BlazeRuleEngine(sparkSession)
+ })
+ }
+
+ override def createConvertToNativeExec(child: SparkPlan): ConvertToNativeBase =
+ ConvertToNativeExec(child)
+
+ override def createNativeAggExec(
+ execMode: AggExecMode,
+ requiredChildDistributionExpressions: Option[Seq[Expression]],
+ groupingExpressions: Seq[NamedExpression],
+ aggregateExpressions: Seq[AggregateExpression],
+ aggregateAttributes: Seq[Attribute],
+ initialInputBufferOffset: Int,
+ child: SparkPlan): NativeAggBase =
+ NativeAggExec(
+ execMode,
+ requiredChildDistributionExpressions,
+ groupingExpressions,
+ aggregateExpressions,
+ aggregateAttributes,
+ initialInputBufferOffset,
+ child)
+
+ override def createNativeBroadcastExchangeExec(
+ mode: BroadcastMode,
+ child: SparkPlan): NativeBroadcastExchangeBase =
+ NativeBroadcastExchangeExec(mode, child)
+
+ override def createNativeBroadcastJoinExec(
+ left: SparkPlan,
+ right: SparkPlan,
+ outputPartitioning: Partitioning,
+ leftKeys: Seq[Expression],
+ rightKeys: Seq[Expression],
+ joinType: JoinType,
+ broadcastSide: BroadcastSide): NativeBroadcastJoinBase =
+ NativeBroadcastJoinExec(
+ left,
+ right,
+ outputPartitioning,
+ leftKeys,
+ rightKeys,
+ joinType,
+ broadcastSide)
+
+ override def createNativeShuffledHashJoinExec(
+ left: SparkPlan,
+ right: SparkPlan,
+ leftKeys: Seq[Expression],
+ rightKeys: Seq[Expression],
+ joinType: JoinType,
+ buildSide: BuildSide,
+ isSkewJoin: Boolean): SparkPlan =
+ NativeShuffledHashJoinExec(left, right, leftKeys, rightKeys, joinType, buildSide)
+
+ override def createNativeSortMergeJoinExec(
+ left: SparkPlan,
+ right: SparkPlan,
+ leftKeys: Seq[Expression],
+ rightKeys: Seq[Expression],
+ joinType: JoinType,
+ isSkewJoin: Boolean): NativeSortMergeJoinBase =
+ NativeSortMergeJoinExec(left, right, leftKeys, rightKeys, joinType)
+
+ override def createNativeExpandExec(
+ projections: Seq[Seq[Expression]],
+ output: Seq[Attribute],
+ child: SparkPlan): NativeExpandBase =
+ NativeExpandExec(projections, output, child)
+
+ override def createNativeFilterExec(condition: Expression, child: SparkPlan): NativeFilterBase =
+ NativeFilterExec(condition, child)
+
+ override def createNativeGenerateExec(
+ generator: Generator,
+ requiredChildOutput: Seq[Attribute],
+ outer: Boolean,
+ generatorOutput: Seq[Attribute],
+ child: SparkPlan): NativeGenerateBase =
+ NativeGenerateExec(generator, requiredChildOutput, outer, generatorOutput, child)
+
+ override def createNativeGlobalLimitExec(limit: Long, child: SparkPlan): NativeGlobalLimitBase =
+ NativeGlobalLimitExec(limit, child)
+
+ override def createNativeLocalLimitExec(limit: Long, child: SparkPlan): NativeLocalLimitBase =
+ NativeLocalLimitExec(limit, child)
+
+ override def createNativeParquetInsertIntoHiveTableExec(
+ cmd: InsertIntoHiveTable,
+ child: SparkPlan): NativeParquetInsertIntoHiveTableBase =
+ NativeParquetInsertIntoHiveTableExec(cmd, child)
+
+ override def createNativeParquetScanExec(
+ basedFileScan: FileSourceScanExec): NativeParquetScanBase =
+ NativeParquetScanExec(basedFileScan)
+
+ override def createNativeOrcScanExec(basedFileScan: FileSourceScanExec): NativeOrcScanBase =
+ NativeOrcScanExec(basedFileScan)
+
+ override def createNativeProjectExec(
+ projectList: Seq[NamedExpression],
+ child: SparkPlan): NativeProjectBase =
+ NativeProjectExec(projectList, child)
+
+ override def createNativeRenameColumnsExec(
+ child: SparkPlan,
+ newColumnNames: Seq[String]): NativeRenameColumnsBase =
+ NativeRenameColumnsExec(child, newColumnNames)
+
+ override def createNativeShuffleExchangeExec(
+ outputPartitioning: Partitioning,
+ child: SparkPlan,
+ shuffleOrigin: Option[Any] = None): NativeShuffleExchangeBase =
+ NativeShuffleExchangeExec(outputPartitioning, child)
+
+ override def createNativeSortExec(
+ sortOrder: Seq[SortOrder],
+ global: Boolean,
+ child: SparkPlan): NativeSortBase =
+ NativeSortExec(sortOrder, global, child)
+
+ override def createNativeTakeOrderedExec(
+ limit: Long,
+ sortOrder: Seq[SortOrder],
+ child: SparkPlan): NativeTakeOrderedBase =
+ NativeTakeOrderedExec(limit, sortOrder, child)
+
+ override def createNativePartialTakeOrderedExec(
+ limit: Long,
+ sortOrder: Seq[SortOrder],
+ child: SparkPlan,
+ metrics: Map[String, SQLMetric]): NativePartialTakeOrderedBase =
+ NativePartialTakeOrderedExec(limit, sortOrder, child, metrics)
+
+ override def createNativeUnionExec(
+ children: Seq[SparkPlan],
+ output: Seq[Attribute]): NativeUnionBase =
+ NativeUnionExec(children, output)
+
+ override def createNativeWindowExec(
+ windowExpression: Seq[NamedExpression],
+ partitionSpec: Seq[Expression],
+ orderSpec: Seq[SortOrder],
+ groupLimit: Option[Int],
+ outputWindowCols: Boolean,
+ child: SparkPlan): NativeWindowBase = {
+ assert(groupLimit.isEmpty, "spark241kwaiae does not support WindowExec with groupLimit")
+ assert(
+ outputWindowCols,
+ "spark241kwaiae does not support WindowExec without outputWindowCols")
+ NativeWindowExec(windowExpression, partitionSpec, orderSpec, child)
+ }
+
+ override def createNativeParquetSinkExec(
+ sparkSession: SparkSession,
+ table: CatalogTable,
+ partition: Map[String, Option[String]],
+ child: SparkPlan,
+ metrics: Map[String, SQLMetric]): NativeParquetSinkBase =
+ NativeParquetSinkExec(sparkSession, table, partition, child, metrics)
+
+ override def getUnderlyingBroadcast(exec: SparkPlan): BroadcastExchangeLike = {
+ exec match {
+ case exec: BroadcastExchangeLike => exec
+ case exec: UnaryExecNode => getUnderlyingBroadcast(exec.child)
+ case exec: BroadcastQueryStageInput => getUnderlyingBroadcast(exec.childStage.child)
+ }
+ }
+
+ override def getRDDShuffleReadFull(rdd: RDD[_]): Boolean = {
+ rdd.shuffleReadFull
+ }
+
+ override def setRDDShuffleReadFull(rdd: RDD[_], shuffleReadFull: Boolean): Unit = {
+ rdd.shuffleReadFull = shuffleReadFull
+ }
+
+ override def getRssPartitionWriter(
+ handle: ShuffleHandle,
+ mapId: Int,
+ metrics: ShuffleWriteMetricsReporter,
+ numPartitions: Int): Option[RssPartitionWriter] = {
+ Some(new RssPartitionWriter(handle, mapId, metrics, numPartitions))
+ }
+
+ override def createFileSegment(
+ file: File,
+ offset: Long,
+ length: Long,
+ numRecords: Long): FileSegment = new FileSegment(file, offset, length, numRecords)
+
+ override def commit(
+ dep: ShuffleDependency[_, _, _],
+ shuffleBlockResolver: IndexShuffleBlockResolver,
+ tempDataFile: File,
+ mapId: Long,
+ partitionLengths: Array[Long],
+ dataSize: Long,
+ context: TaskContext): MapStatus = {
+
+ // commit
+ shuffleBlockResolver.writeIndexFileAndCommit(
+ dep.shuffleId,
+ mapId.toInt,
+ partitionLengths,
+ tempDataFile)
+
+ // shuffle write on hdfs
+ val blockManager = SparkEnv.get.blockManager
+ val handle = dep.shuffleHandle.asInstanceOf[BaseShuffleHandle[_, _, _]]
+ val output = shuffleBlockResolver.getDataFile(dep.shuffleId, mapId.toInt)
+ val mapInfo = new MapInfo(partitionLengths, partitionLengths.map(_ => 0L))
+ val totalPartitionLength = mapInfo.lengths.sum
+ var hasExternalData = false
+
+ if (ExternalBlockStoreUtils.writeRemoteEnabled(
+ SparkEnv.get.conf,
+ totalPartitionLength,
+ handle.numMaps)) {
+ if (dataSize != 0) {
+ logInfo(s"KwaiShuffle: Start to write remote, shuffle block id ${ShuffleDataBlockId(
+ dep.shuffleId,
+ mapId.toInt,
+ 0).name}, file path is ${output.getPath}/${output.getName}")
+ context.taskMetrics.externalMetrics.writeRemoteShuffle.setValue(1L)
+ val dataBlockId = new ShuffleDataBlockId(
+ dep.shuffleId,
+ mapId.toInt,
+ IndexShuffleBlockResolver.NOOP_REDUCE_ID)
+ val indexFile = shuffleBlockResolver.getIndexFile(dep.shuffleId, mapId.toInt)
+ blockManager.externalBlockStore.externalBlockManager.get.writeExternalShuffleFile(
+ dep.shuffleId,
+ mapId.toInt,
+ dataBlockId,
+ indexFile,
+ output)
+ hasExternalData = true
+ }
+ }
+
+ val mapStatus =
+ MapStatus.apply(SparkEnv.get.blockManager.shuffleServerId, partitionLengths)
+ mapStatus.hasExternal = hasExternalData
+ mapStatus
+ }
+
+ override def getMapStatus(
+ shuffleServerId: BlockManagerId,
+ partitionLengthMap: Array[Long],
+ mapId: Long): MapStatus = {
+ val mapStatus = MapStatus.apply(shuffleServerId, partitionLengthMap)
+ mapStatus.hasExternal = false
+ mapStatus
+ }
+
+ override def getShuffleWriteExec(
+ input: pb.PhysicalPlanNode,
+ nativeOutputPartitioning: pb.PhysicalRepartition.Builder): pb.PhysicalPlanNode = {
+ if (SparkEnv.get.conf.get("spark.shuffle.manager", "sort").equals("unify")) {
+ pb.PhysicalPlanNode
+ .newBuilder()
+ .setRssShuffleWriter(
+ pb.RssShuffleWriterExecNode
+ .newBuilder()
+ .setInput(input)
+ .setOutputPartitioning(nativeOutputPartitioning)
+ .buildPartial()
+ ) // shuffleId is not set at the moment, will be set in ShuffleWriteProcessor
+ .build()
+ } else {
+ pb.PhysicalPlanNode
+ .newBuilder()
+ .setShuffleWriter(
+ pb.ShuffleWriterExecNode
+ .newBuilder()
+ .setInput(input)
+ .setOutputPartitioning(nativeOutputPartitioning)
+ .buildPartial()
+ ) // shuffleId is not set at the moment, will be set in ShuffleWriteProcessor
+ .build()
+ }
+ }
+
+ override def isNative(plan: SparkPlan): Boolean = {
+ plan match {
+ case _: NativeSupports => true
+ case plan: QueryStageInput => isNative(plan.childStage)
+ case plan: QueryStage => isNative(plan.child)
+ case plan: ReusedExchangeExec => isNative(plan.child)
+ case _ => false
+ }
+ }
+
+ override def getUnderlyingNativePlan(plan: SparkPlan): NativeSupports = {
+ plan match {
+ case plan: NativeSupports => plan
+ case plan: QueryStageInput => getUnderlyingNativePlan(plan.childStage)
+ case plan: QueryStage => getUnderlyingNativePlan(plan.child)
+ case plan: ReusedExchangeExec => getUnderlyingNativePlan(plan.child)
+ case _ => throw new RuntimeException("unreachable: plan is not native")
+ }
+ }
+
+ override def executeNative(plan: SparkPlan): NativeRDD =
+ plan match {
+ case plan: NativeSupports => plan.executeNative()
+ case plan: ShuffleQueryStageInput => executeNativeCustomShuffleReader(plan)
+ case plan: SkewedShuffleQueryStageInput => executeNativeCustomShuffleReader(plan)
+ case plan: BroadcastQueryStageInput => executeNative(plan.childStage)
+ case plan: QueryStage => executeNative(plan.child)
+ case plan: ReusedExchangeExec => executeNative(plan.child)
+ case _ => throw new SparkException(s"Underlying plan is not NativeSupports: ${plan}")
+ }
+
+ override def isQueryStageInput(plan: SparkPlan): Boolean =
+ plan.isInstanceOf[QueryStageInput]
+
+ override def isShuffleQueryStageInput(plan: SparkPlan): Boolean = {
+ plan.isInstanceOf[ShuffleQueryStageInput] || plan.isInstanceOf[SkewedShuffleQueryStageInput]
+ }
+
+ override def getChildStage(plan: SparkPlan): SparkPlan =
+ plan.asInstanceOf[QueryStageInput].childStage
+
+ override def setLogicalLink(exec: SparkPlan, basedExec: SparkPlan): SparkPlan =
+ exec
+
+ override def simpleStringWithNodeId(plan: SparkPlan): String = plan.simpleString
+
+ override def convertMoreExprWithFallback(
+ e: Expression,
+ isPruningExpr: Boolean,
+ fallback: Expression => pb.PhysicalExprNode): Option[pb.PhysicalExprNode] = {
+ e match {
+ case StringSplit(str, pat @ Literal(_, StringType))
+ // native StringSplit implementation does not support regex, so only most frequently
+ // used cases without regex are supported
+ if Seq(",", ", ", ":", ";", "#", "@", "_", "-", "\\|", "\\.").contains(pat.value) =>
+ val nativePat = pat.value match {
+ case "\\|" => "|"
+ case "\\." => "."
+ case other => other
+ }
+ Some(
+ pb.PhysicalExprNode
+ .newBuilder()
+ .setScalarFunction(
+ pb.PhysicalScalarFunctionNode
+ .newBuilder()
+ .setFun(pb.ScalarFunction.SparkExtFunctions)
+ .setName("StringSplit")
+ .addArgs(NativeConverters.convertExpr(str))
+ .addArgs(NativeConverters.convertExpr(Literal(nativePat)))
+ .setReturnType(NativeConverters.convertDataType(StringType)))
+ .build())
+
+ case PromotePrecision(_1) =>
+ throw new NotImplementedError("PromotePrecision is disabled")
+
+ case _ => None
+ }
+ }
+
+ override def getLikeEscapeChar(expr: Expression): Char = '\\'
+
+ override def convertMoreAggregateExpr(e: AggregateExpression): Option[pb.PhysicalExprNode] = {
+ assert(getAggregateExpressionFilter(e).isEmpty)
+ val aggBuilder = pb.PhysicalAggExprNode.newBuilder()
+
+ e.aggregateFunction match {
+ case First(child, ignoresNullExpr) =>
+ aggBuilder.setAggFunction(if (ignoresNullExpr.eval().asInstanceOf[Boolean]) {
+ pb.AggFunction.FIRST_IGNORES_NULL
+ } else {
+ pb.AggFunction.FIRST
+ })
+ aggBuilder.addChildren(NativeConverters.convertExpr(child))
+ Some(pb.PhysicalExprNode.newBuilder().setAggExpr(aggBuilder).build())
+
+ case _ => None
+ }
+ }
+
+ override def getAggregateExpressionFilter(expr: Expression): Option[Expression] = None
+
+ private def executeNativeCustomShuffleReader(exec: SparkPlan): NativeRDD = {
+ val shuffledRDD = exec.execute()
+ val dependency = shuffledRDD.getClass
+ .getMethod("dependency")
+ .invoke(shuffledRDD)
+ .asInstanceOf[ShuffleDependency[_, _, _]]
+ val shuffleHandle = dependency.shuffleHandle
+
+ val shuffleExec = exec
+ .asInstanceOf[QueryStageInput]
+ .childStage
+ .child
+ .asInstanceOf[NativeShuffleExchangeExec]
+ val metrics = shuffleExec.metrics
+ val nativeMetrics = MetricNode(
+ Map(),
+ Nil,
+ Some({
+ case ("output_rows", v) =>
+ val taskMetrics = TaskContext.get.taskMetrics()
+ val shuffleReadMetrics = taskMetrics.createTempShuffleReadMetrics()
+ new SQLShuffleReadMetricsReporter(shuffleReadMetrics, metrics).incRecordsRead(v)
+ taskMetrics.mergeShuffleReadMetrics()
+ case ("elapsed_compute", v) => metrics("shuffle_read_total_time") += v
+ case _ =>
+ }))
+ val nativeSchema = shuffleExec.nativeSchema
+ val partitionClsName = shuffledRDD.getClass.getSimpleName
+ val rebalancedPartitions = shuffledRDD match {
+ case rdd: LocalShuffledRowRDD =>
+ Some(
+ FieldUtils
+ .readDeclaredField(
+ rdd,
+ "org$apache$spark$sql$execution$adaptive$LocalShuffledRowRDD$$rebalancedPartitions",
+ true)
+ .asInstanceOf[Array[_]])
+ case _ =>
+ None
+ }
+
+ new NativeRDD(
+ shuffledRDD.sparkContext,
+ nativeMetrics,
+ shuffledRDD.partitions,
+ shuffledRDD.partitioner,
+ new OneToOneDependency(shuffledRDD) :: Nil,
+ shuffledRDD.shuffleReadFull,
+ (partition, taskContext) => {
+ val shuffleReadMetrics = taskContext.taskMetrics().createTempShuffleReadMetrics()
+ val metricsReporter =
+ new SQLShuffleReadMetricsReporter(shuffleReadMetrics, metrics)
+
+ val classOfShuffledRowRDDPartition =
+ Utils.classForName("org.apache.spark.sql.execution.ShuffledRowRDDPartition")
+ val classOfAdaptiveShuffledRowRDDPartition =
+ Utils.classForName(
+ "org.apache.spark.sql.execution.adaptive.AdaptiveShuffledRowRDDPartition")
+
+ val readers: Iterator[ShuffleReader[_, _]] = shuffledRDD match {
+ case rdd: LocalShuffledRowRDD =>
+ ShimsImpl
+ .getLocalShuffleReaders(
+ taskContext,
+ rdd,
+ dependency,
+ rebalancedPartitions.get,
+ partition)
+ .toIterator
+ case _ =>
+ partition match {
+ case p if classOfShuffledRowRDDPartition.isInstance(p) =>
+ val clz = classOfShuffledRowRDDPartition
+ val startPreShufflePartitionIndex =
+ clz.getMethod("startPreShufflePartitionIndex").invoke(p).asInstanceOf[Int]
+ val endPreShufflePartitionIndex =
+ clz.getMethod("endPreShufflePartitionIndex").invoke(p).asInstanceOf[Int]
+
+ Iterator.single(
+ SparkEnv.get.shuffleManager
+ .getReader(
+ shuffleHandle,
+ startPreShufflePartitionIndex,
+ endPreShufflePartitionIndex,
+ taskContext,
+ metricsReporter))
+
+ case p if classOfAdaptiveShuffledRowRDDPartition.isInstance(p) =>
+ val clz = classOfAdaptiveShuffledRowRDDPartition
+ val preShufflePartitionIndex =
+ clz.getMethod("preShufflePartitionIndex").invoke(p).asInstanceOf[Int]
+ val startMapId = clz.getMethod("startMapId").invoke(p).asInstanceOf[Int]
+ val endMapId = clz.getMethod("endMapId").invoke(p).asInstanceOf[Int]
+
+ Iterator.single(
+ SparkEnv.get.shuffleManager
+ .getReader(
+ shuffleHandle,
+ preShufflePartitionIndex,
+ preShufflePartitionIndex + 1,
+ taskContext,
+ metricsReporter,
+ startMapId,
+ endMapId))
+ case p =>
+ Iterator.single(
+ SparkEnv.get.shuffleManager
+ .getReader(shuffleHandle, p.index, p.index + 1, taskContext, metricsReporter))
+ }
+ }
+
+ // store fetch iterator in jni resource before native compute
+ val jniResourceId = s"NativeShuffleReadExec:${UUID.randomUUID().toString}"
+ JniBridge.resourcesMap.put(
+ jniResourceId,
+ () => {
+ CompletionIterator[Object, Iterator[Object]](
+ readers.flatMap(_.asInstanceOf[BlazeBlockStoreShuffleReaderBase[_, _]].readIpc()),
+ taskContext.taskMetrics().mergeShuffleReadMetrics(true))
+ })
+
+ pb.PhysicalPlanNode
+ .newBuilder()
+ .setIpcReader(
+ pb.IpcReaderExecNode
+ .newBuilder()
+ .setSchema(nativeSchema)
+ .setNumPartitions(shuffledRDD.getNumPartitions)
+ .setIpcProviderResourceId(jniResourceId)
+ .build())
+ .build()
+ },
+ friendlyName = s"NativeRDD.ShuffleRead [$partitionClsName]")
+ }
+
+ override def convertMoreSparkPlan(exec: SparkPlan): Option[SparkPlan] = {
+ exec match {
+ case _: ShuffleQueryStageInput | _: SkewedShuffleQueryStageInput | _: ReusedExchangeExec
+ if Shims.get.isNative(exec) =>
+ Some(ForceNativeExecutionWrapper(BlazeConverters.addRenameColumnsExec(exec)))
+
+ case _ =>
+ None
+ }
+ }
+
+ override def getSqlContext(sparkPlan: SparkPlan): SQLContext = sparkPlan.sqlContext
+
+ override def createNativeExprWrapper(
+ nativeExpr: PhysicalExprNode,
+ dataType: DataType,
+ nullable: Boolean): Expression = {
+ case class NativeExprWrapper(
+ nativeExpr: PhysicalExprNode,
+ override val dataType: DataType,
+ override val nullable: Boolean)
+ extends NativeExprWrapperBase(nativeExpr, dataType, nullable)
+ NativeExprWrapper(nativeExpr, dataType, nullable)
+ }
+
+ override def postTransform(plan: SparkPlan, sc: SparkContext): Unit = {
+ if (!sc.conf.getBoolean("spark.blaze.enable.addListener", defaultValue = false)) {
+ sc.conf.set("spark.blaze.enable.addListener", "true")
+ val statusStore: SQLAppStatusStore = {
+ val kvStore = sc.statusStore.store.asInstanceOf[ElementTrackingStore]
+ val listener = new SQLAppStatusListener(sc.conf, kvStore, live = true)
+ sc.listenerBus.addToStatusQueue(listener)
+ val statusStore = new SQLAppStatusStore(kvStore, Some(listener))
+ statusStore
+ }
+ sc.listenerBus.addToQueue(new BlazePlanListener(sc, statusStore), "BlazeOperatorMetrics")
+ }
+ }
+
+ def getMinPartitionNum(sparkSession: org.apache.spark.sql.SparkSession): Int =
+ sparkSession.sparkContext.defaultParallelism
+
+ def getPartitionedFile(
+ partitionValues: org.apache.spark.sql.catalyst.InternalRow,
+ filePath: String,
+ offset: Long,
+ size: Long): org.apache.spark.sql.execution.datasources.PartitionedFile =
+ throw new UnsupportedOperationException("not supported in spark241kwaiae")
+}
+
+object ShimsImpl extends Logging {
+
+ // copied from spark LocalShuffleRowRDD
+ def getLocalShuffleReaders(
+ context: TaskContext,
+ rdd: LocalShuffledRowRDD,
+ dependency: ShuffleDependency[_, _, _],
+ rebalancedPartitions: Array[_],
+ split: Partition): Iterable[ShuffleReader[_, _]] = {
+
+ val shuffledRowPartition = split.asInstanceOf[ShuffledRDDPartition]
+ val mapId = shuffledRowPartition.index
+ val metrics = context.taskMetrics().createTempShuffleReadMetrics()
+
+ if (rebalancedPartitions.nonEmpty) {
+ val rebalancedPartition = rebalancedPartitions(mapId)
+ val segments =
+ MethodUtils.invokeMethod(rebalancedPartition, true, "segments").asInstanceOf[Seq[_]]
+ val groupedRebalancedPartitions = segments.groupBy(segment =>
+ MethodUtils.invokeMethod(segment, true, "partId").asInstanceOf[Int])
+
+ groupedRebalancedPartitions.map { case (partId, partitionStatistics) =>
+ SparkEnv.get.shuffleManager
+ .asInstanceOf[UnifiedShuffleManager]
+ .getRebalancedReader(
+ dependency.shuffleHandle,
+ partId,
+ context,
+ metrics,
+ Seq({
+ val fileInfo = new java.util.HashMap[String, java.lang.Long]()
+ partitionStatistics.foreach { segment =>
+ val fileName =
+ MethodUtils.invokeMethod(segment, true, "fileName").asInstanceOf[String]
+ val partId = MethodUtils.invokeMethod(segment, true, "partId").asInstanceOf[Int]
+ fileInfo.put(fileName, partId)
+ }
+ PartitionStatistics
+ .builder()
+ .stageId(dependency.shuffleId)
+ .partId(partId)
+ .fileInfo(fileInfo)
+ .build()
+ }))
+ }
+
+ } else {
+ rdd.partitionStartIndices
+ .zip(rdd.partitionEndIndices)
+ .map { case (start, end) =>
+ logInfo(s"Create local shuffle reader mapId $mapId, partition range $start-$end")
+ SparkEnv.get.shuffleManager
+ .getReader(dependency.shuffleHandle, start, end, context, metrics, mapId, mapId + 1)
+ }
+ }
+ }
+}
+
+case class BlazeRuleEngine(sparkSession: SparkSession) extends Rule[LogicalPlan] with Logging {
+ import BlazeSparkSessionExtension._
+
+ override def apply(plan: LogicalPlan): LogicalPlan = {
+ plan.foreachUp {
+ case p @ LogicalRelation(fsRelation: HadoopFsRelation, _, _, _) =>
+ // non parquet table rule
+ if (!(
+ BlazeConverters.enableScanParquet && fsRelation.fileFormat.isInstanceOf[ParquetFileFormat]
+ || BlazeConverters.enableScanOrc && fsRelation.fileFormat.isInstanceOf[OrcFileFormat]
+ )) {
+ turnOffBlazeWithReason(p.conf, BlazeMissPatterns.NonParquetFormat)
+ }
+
+ // read encrypted table rule
+ val readEncryptedTableEnable = sparkSession.sparkContext.conf
+ .getBoolean("spark.hive.exist.read.encrypted.table", defaultValue = false)
+ if (readEncryptedTableEnable) {
+ turnOffBlazeWithReason(p.conf, BlazeMissPatterns.ReadEncryptedTable)
+ }
+
+ // skip scan dp_dd.*** table because parquet statics don't get min/max info
+ if (p.catalogTable.map(_.identifier.unquotedString).getOrElse("").contains("dp_dd")) {
+ turnOffBlazeWithReason(p.conf, BlazeMissPatterns.ReadHbaseTable)
+ }
+
+ // skip scan offline_attribution_mover_v2
+ // issue: https://team.corp.kuaishou.com/task/T6538018
+ if (p.catalogTable
+ .map(_.identifier.unquotedString)
+ .getOrElse("")
+ .contains("offline_attribution_mover_v2")) {
+ turnOffBlazeWithReason(p.conf, BlazeMissPatterns.ReadBlacklistedTable)
+ }
+
+ case h: HiveTableRelation =>
+ turnOffBlazeWithReason(h.conf, BlazeMissPatterns.NonParquetFormat)
+
+ case _ =>
+ }
+ plan
+ }
+
+ private def turnOffBlazeWithReason(planConf: SQLConf, blazeMissPattern: String): Unit = {
+ planConf.setConf(blazeEnabledKey, false)
+ sparkSession.sparkContext.conf
+ .set(BlazeRuleEngine.blazeMissPatterns, blazeMissPattern)
+ }
+
+ object BlazeMissPatterns extends Enumeration {
+ val NonParquetFormat = "NonParquetFormat"
+ val ReadEncryptedTable = "ReadEncryptedTable"
+ val ReadHbaseTable = "ReadHbaseTable"
+ val ReadBlacklistedTable = "ReadBlacklistedTable"
+ }
+}
+
+object BlazeRuleEngine {
+ lazy val blazeMissPatterns: OptionalConfigEntry[String] = SQLConf
+ .buildConf("spark.blaze.blazeMissPatterns")
+ .stringConf
+ .createOptional
+}
+
+case class ForceNativeExecutionWrapper(override val child: SparkPlan)
+ extends ForceNativeExecutionWrapperBase(child) {
+
+ override def withNewChildren(newChildren: Seq[SparkPlan]): SparkPlan =
+ copy(child = newChildren.head)
+
+ override protected def doCanonicalize(): SparkPlan = child.canonicalized
+}
diff --git a/spark-extension-shims-spark241kwaiae/src/main/scala/org/apache/spark/sql/blaze/kwai/BlazeOperatorMetricsCollector.scala b/spark-extension-shims-spark241kwaiae/src/main/scala/org/apache/spark/sql/blaze/kwai/BlazeOperatorMetricsCollector.scala
new file mode 100644
index 00000000..186ecc24
--- /dev/null
+++ b/spark-extension-shims-spark241kwaiae/src/main/scala/org/apache/spark/sql/blaze/kwai/BlazeOperatorMetricsCollector.scala
@@ -0,0 +1,130 @@
+/*
+ * Copyright 2022 The Blaze Authors
+ *
+ * Licensed 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.blaze.kwai
+
+import scala.beans.BeanProperty
+import scala.util.Try
+
+import org.apache.spark.scheduler._
+import org.apache.spark.sql.execution.SparkPlan
+import org.apache.spark.{KwaiSparkBasicMetrics, SparkContext, SparkEnv}
+
+import com.fasterxml.jackson.databind.ObjectMapper
+import com.fasterxml.jackson.databind.node.ObjectNode
+import org.apache.spark.internal.Logging
+import org.apache.spark.sql.blaze.kwai.BlazeOperatorMetricsCollector.planStore
+import org.codehaus.jackson.annotate.JsonProperty
+import scala.collection.mutable.{ArrayBuffer, ListBuffer}
+
+import org.apache.spark.sql.blaze.Shims
+
+class BlazeOperatorMetricsCollector extends Logging {
+
+ private val kafkaStageMetricsBroker =
+ SparkEnv.get.conf.get(
+ "spark.kwai.blazeOperatorMetrics.broker",
+ defaultValue =
+ "dataarch-bjdy-rs7110.idczw.hb1.kwaidc.com:9092,dataarch-bjdy-rs7111.idczw.hb1.kwaidc.com:9092,dataarch-bjdy-rs7112.idczw.hb1.kwaidc.com:9092,dataarch-bjdy-rs7113.idczw.hb1.kwaidc.com:9092,dataarch-bjdy-rs7114.idczw.hb1.kwaidc.com:9092")
+ private val kafkaStageMetricsTopic =
+ SparkEnv.get.conf.get(
+ "spark.kwai.blazeOperatorMetrics.topic",
+ defaultValue = "dataarch_blaze_operator_metrics")
+
+ private val producer: Option[KafkaClient] = Try(
+ KafkaClient(kafkaStageMetricsBroker, null, kafkaStageMetricsTopic)).toOption
+
+ private val objectMapper: ObjectMapper = new ObjectMapper()
+
+ def sendOperatorMetrics(sc: SparkContext): Unit = {
+ try {
+ val startTime = System.currentTimeMillis()
+ val operatorList = new ListBuffer[BlazeOperatorInfo]()
+
+ planStore.zipWithIndex.foreach { case (plan, index) =>
+ val metricObject = objectMapper.createObjectNode()
+ plan.metrics.foreach { case (key, value) =>
+ metricObject.put(key, if (value.value < 0) 0L else value.value)
+ }
+ operatorList.append(
+ BlazeOperatorInfo(
+ plan.getClass.getSimpleName,
+ index,
+ Shims.get.isNative(plan),
+ metricObject))
+ }
+ producer.foreach(
+ _.send(
+ key = sc.conf.getAppId,
+ objectMapper.writeValueAsString(new BlazeApplicationInfo(sc, operatorList.toArray))))
+ planStore.clear()
+ logInfo(
+ s"send app operator metrics to kafka succ after ${System.currentTimeMillis() - startTime}ms")
+ } catch {
+ case e: Exception =>
+ logInfo(s"Blaze sendOperatorMetrics error: ${e.getMessage}")
+ producer.foreach(_.flush())
+ producer.foreach(_.close())
+ }
+ }
+}
+
+object BlazeOperatorMetricsCollector {
+ val isBlazeOperatorMetricsEnabled: Boolean =
+ SparkEnv.get.conf.getBoolean("spark.blaze.enable.operatorMetrics", defaultValue = true)
+
+ val instance: Option[BlazeOperatorMetricsCollector] = {
+ if (isBlazeOperatorMetricsEnabled) {
+ Some(new BlazeOperatorMetricsCollector)
+ } else {
+ None
+ }
+ }
+ var planStore: ArrayBuffer[SparkPlan] = new ArrayBuffer[SparkPlan]()
+}
+
+class BlazeApplicationInfo(sc: SparkContext, info: Array[BlazeOperatorInfo])
+ extends KwaiSparkBasicMetrics(sc) {
+ @BeanProperty
+ @JsonProperty("blazeJarVersion")
+ val blazeJarVersion: String = sc.conf.get("spark.blaze.jar", "default")
+ @BeanProperty
+ @JsonProperty("operatorsInfo")
+ val operatorsInfo: Array[BlazeOperatorInfo] = info
+}
+
+case class BlazeOperatorInfo(
+ @BeanProperty
+ @JsonProperty("execName")
+ execName: String,
+ @BeanProperty
+ @JsonProperty("execId")
+ execId: Int,
+ @BeanProperty
+ @JsonProperty("isNative")
+ isNative: Boolean,
+ @BeanProperty
+ @JsonProperty("execMetric")
+ execMetric: ObjectNode) {}
+
+class OperatorMetricsListener(sc: SparkContext) extends SparkListener with Logging {
+
+ override def onApplicationEnd(applicationEnd: SparkListenerApplicationEnd): Unit = {
+ BlazeOperatorMetricsCollector.instance.foreach { collector =>
+ collector.sendOperatorMetrics(sc)
+ }
+ sc.listenerBus.removeListener(this)
+ }
+}
diff --git a/spark-extension-shims-spark241kwaiae/src/main/scala/org/apache/spark/sql/blaze/kwai/KwaiPrivilegedHDFSBlockManager.scala b/spark-extension-shims-spark241kwaiae/src/main/scala/org/apache/spark/sql/blaze/kwai/KwaiPrivilegedHDFSBlockManager.scala
new file mode 100644
index 00000000..eff29d69
--- /dev/null
+++ b/spark-extension-shims-spark241kwaiae/src/main/scala/org/apache/spark/sql/blaze/kwai/KwaiPrivilegedHDFSBlockManager.scala
@@ -0,0 +1,51 @@
+/*
+ * Copyright 2022 The Blaze Authors
+ *
+ * Licensed 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.blaze.kwai
+
+import java.security.PrivilegedExceptionAction
+
+import org.apache.spark.sql.blaze.NativeHelper
+import org.apache.spark.storage.BlockManager
+import org.apache.spark.storage.HDFSBlockManager
+import org.apache.spark.SparkConf
+import org.apache.spark.internal.Logging
+import org.apache.spark.storage.ExternalBlockStore
+
+class KwaiPrivilegedHDFSBlockManager extends HDFSBlockManager {
+ override def init(blockManager: BlockManager): Unit = {
+ val initHDFSBlockManager = () => {
+ super.init(blockManager)
+ }
+
+ NativeHelper.currentUser.doAs(new PrivilegedExceptionAction[Unit] {
+ override def run(): Unit = initHDFSBlockManager()
+ })
+ }
+}
+
+object KwaiPrivilegedHDFSBlockManager extends Logging {
+ def setup(conf: SparkConf): Unit = {
+ val externalBlockManagerClassName = conf
+ .getOption(ExternalBlockStore.BLOCK_MANAGER_NAME)
+ .getOrElse(ExternalBlockStore.DEFAULT_BLOCK_MANAGER_NAME)
+ val thisClassName = classOf[KwaiPrivilegedHDFSBlockManager].getName
+
+ if (externalBlockManagerClassName == classOf[HDFSBlockManager].getName) {
+ logInfo(s"replacing external block manager to $thisClassName")
+ conf.set(ExternalBlockStore.BLOCK_MANAGER_NAME, thisClassName)
+ }
+ }
+}
diff --git a/spark-extension-shims-spark3/src/main/scala/org/apache/spark/shuffle/uniffle/RssShuffleHandleWrapper.scala b/spark-extension-shims-spark241kwaiae/src/main/scala/org/apache/spark/sql/execution/blaze/plan/ConvertToNativeExec.scala
similarity index 64%
copy from spark-extension-shims-spark3/src/main/scala/org/apache/spark/shuffle/uniffle/RssShuffleHandleWrapper.scala
copy to spark-extension-shims-spark241kwaiae/src/main/scala/org/apache/spark/sql/execution/blaze/plan/ConvertToNativeExec.scala
index e8486b0a..9c77bf70 100644
--- a/spark-extension-shims-spark3/src/main/scala/org/apache/spark/shuffle/uniffle/RssShuffleHandleWrapper.scala
+++ b/spark-extension-shims-spark241kwaiae/src/main/scala/org/apache/spark/sql/execution/blaze/plan/ConvertToNativeExec.scala
@@ -13,11 +13,11 @@
* See the License for the specific language governing permissions and
* limitations under the License.
*/
-package org.apache.spark.shuffle.uniffle
+package org.apache.spark.sql.execution.blaze.plan
-import org.apache.spark.shuffle.{BaseShuffleHandle, RssShuffleHandle}
+import org.apache.spark.sql.execution.SparkPlan
-class RssShuffleHandleWrapper[K, V, C](val rssShuffleHandleInfo: RssShuffleHandle[K, V, C])
- extends BaseShuffleHandle[K, V, C](
- rssShuffleHandleInfo.getShuffleId,
- rssShuffleHandleInfo.getDependency) {}
+case class ConvertToNativeExec(override val child: SparkPlan) extends ConvertToNativeBase(child) {
+ override def withNewChildren(newChildren: Seq[SparkPlan]): SparkPlan =
+ copy(child = newChildren.head)
+}
diff --git a/spark-extension-shims-spark241kwaiae/src/main/scala/org/apache/spark/sql/execution/blaze/plan/NativeAggExec.scala b/spark-extension-shims-spark241kwaiae/src/main/scala/org/apache/spark/sql/execution/blaze/plan/NativeAggExec.scala
new file mode 100644
index 00000000..9cffd075
--- /dev/null
+++ b/spark-extension-shims-spark241kwaiae/src/main/scala/org/apache/spark/sql/execution/blaze/plan/NativeAggExec.scala
@@ -0,0 +1,44 @@
+/*
+ * Copyright 2022 The Blaze Authors
+ *
+ * Licensed 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.execution.blaze.plan
+
+import org.apache.spark.sql.catalyst.expressions.Attribute
+import org.apache.spark.sql.catalyst.expressions.aggregate.AggregateExpression
+import org.apache.spark.sql.catalyst.expressions.Expression
+import org.apache.spark.sql.catalyst.expressions.NamedExpression
+import org.apache.spark.sql.execution.SparkPlan
+import org.apache.spark.sql.execution.blaze.plan.NativeAggBase.AggExecMode
+
+case class NativeAggExec(
+ execMode: AggExecMode,
+ requiredChildDistributionExpressions: Option[Seq[Expression]],
+ groupingExpressions: Seq[NamedExpression],
+ aggregateExpressions: Seq[AggregateExpression],
+ aggregateAttributes: Seq[Attribute],
+ initialInputBufferOffset: Int,
+ override val child: SparkPlan)
+ extends NativeAggBase(
+ execMode,
+ requiredChildDistributionExpressions,
+ groupingExpressions,
+ aggregateExpressions,
+ aggregateAttributes,
+ initialInputBufferOffset,
+ child) {
+
+ override def withNewChildren(newChildren: Seq[SparkPlan]): SparkPlan =
+ copy(child = newChildren.head)
+}
diff --git a/spark-extension-shims-spark241kwaiae/src/main/scala/org/apache/spark/sql/execution/blaze/plan/NativeBroadcastExchangeExec.scala b/spark-extension-shims-spark241kwaiae/src/main/scala/org/apache/spark/sql/execution/blaze/plan/NativeBroadcastExchangeExec.scala
new file mode 100644
index 00000000..118e611b
--- /dev/null
+++ b/spark-extension-shims-spark241kwaiae/src/main/scala/org/apache/spark/sql/execution/blaze/plan/NativeBroadcastExchangeExec.scala
@@ -0,0 +1,53 @@
+/*
+ * Copyright 2022 The Blaze Authors
+ *
+ * Licensed 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.execution.blaze.plan
+
+import java.util.UUID
+
+import org.apache.commons.lang3.reflect.MethodUtils
+import org.apache.spark.broadcast.Broadcast
+import org.apache.spark.sql.catalyst.plans.physical.BroadcastMode
+import org.apache.spark.sql.execution.SparkPlan
+
+case class NativeBroadcastExchangeExec(mode: BroadcastMode, override val child: SparkPlan)
+ extends NativeBroadcastExchangeBase(mode, child) {
+
+ private val runId: UUID = UUID.randomUUID()
+ override def getRunId: UUID = runId
+
+ override def doExecuteBroadcast[T](): Broadcast[T] = {
+ val broadcast = super.doExecuteBroadcast[Any]()
+ val v = broadcast.value
+
+ val dummyBroadcasted = new Broadcast[Any](-1) {
+ override protected def getValue(): Any = (v, 0, 0, 0, 0) // modified in spark241kwaiae
+
+ override def setValue(v: () => (Any, Long => Any)): Any =
+ throw new UnsupportedOperationException("dummyBroadcasted.setValue")
+
+ override protected def doUnpersist(blocking: Boolean): Unit = {
+ MethodUtils.invokeMethod(broadcast, true, "doUnpersist", Array(blocking))
+ }
+ override protected def doDestroy(blocking: Boolean): Unit = {
+ MethodUtils.invokeMethod(broadcast, true, "doDestroy", Array(blocking))
+ }
+ }
+ dummyBroadcasted.asInstanceOf[Broadcast[T]]
+ }
+
+ override def withNewChildren(newChildren: Seq[SparkPlan]): SparkPlan =
+ copy(child = newChildren.head)
+}
diff --git a/spark-extension-shims-spark3/src/main/scala/org/apache/spark/shuffle/uniffle/RssShuffleHandleWrapper.scala b/spark-extension-shims-spark241kwaiae/src/main/scala/org/apache/spark/sql/execution/blaze/plan/NativeExpandExec.scala
similarity index 53%
copy from spark-extension-shims-spark3/src/main/scala/org/apache/spark/shuffle/uniffle/RssShuffleHandleWrapper.scala
copy to spark-extension-shims-spark241kwaiae/src/main/scala/org/apache/spark/sql/execution/blaze/plan/NativeExpandExec.scala
index e8486b0a..5d6d5060 100644
--- a/spark-extension-shims-spark3/src/main/scala/org/apache/spark/shuffle/uniffle/RssShuffleHandleWrapper.scala
+++ b/spark-extension-shims-spark241kwaiae/src/main/scala/org/apache/spark/sql/execution/blaze/plan/NativeExpandExec.scala
@@ -13,11 +13,18 @@
* See the License for the specific language governing permissions and
* limitations under the License.
*/
-package org.apache.spark.shuffle.uniffle
+package org.apache.spark.sql.execution.blaze.plan
-import org.apache.spark.shuffle.{BaseShuffleHandle, RssShuffleHandle}
+import org.apache.spark.sql.catalyst.expressions.Attribute
+import org.apache.spark.sql.catalyst.expressions.Expression
+import org.apache.spark.sql.execution.SparkPlan
-class RssShuffleHandleWrapper[K, V, C](val rssShuffleHandleInfo: RssShuffleHandle[K, V, C])
- extends BaseShuffleHandle[K, V, C](
- rssShuffleHandleInfo.getShuffleId,
- rssShuffleHandleInfo.getDependency) {}
+case class NativeExpandExec(
+ projections: Seq[Seq[Expression]],
+ override val output: Seq[Attribute],
+ override val child: SparkPlan)
+ extends NativeExpandBase(projections, output, child) {
+
+ override def withNewChildren(newChildren: Seq[SparkPlan]): SparkPlan =
+ copy(child = newChildren.head)
+}
diff --git a/spark-extension-shims-spark3/src/main/scala/org/apache/spark/shuffle/uniffle/RssShuffleHandleWrapper.scala b/spark-extension-shims-spark241kwaiae/src/main/scala/org/apache/spark/sql/execution/blaze/plan/NativeFilterExec.scala
similarity index 60%
copy from spark-extension-shims-spark3/src/main/scala/org/apache/spark/shuffle/uniffle/RssShuffleHandleWrapper.scala
copy to spark-extension-shims-spark241kwaiae/src/main/scala/org/apache/spark/sql/execution/blaze/plan/NativeFilterExec.scala
index e8486b0a..9d343a8c 100644
--- a/spark-extension-shims-spark3/src/main/scala/org/apache/spark/shuffle/uniffle/RssShuffleHandleWrapper.scala
+++ b/spark-extension-shims-spark241kwaiae/src/main/scala/org/apache/spark/sql/execution/blaze/plan/NativeFilterExec.scala
@@ -13,11 +13,14 @@
* See the License for the specific language governing permissions and
* limitations under the License.
*/
-package org.apache.spark.shuffle.uniffle
+package org.apache.spark.sql.execution.blaze.plan
-import org.apache.spark.shuffle.{BaseShuffleHandle, RssShuffleHandle}
+import org.apache.spark.sql.catalyst.expressions.Expression
+import org.apache.spark.sql.execution.SparkPlan
-class RssShuffleHandleWrapper[K, V, C](val rssShuffleHandleInfo: RssShuffleHandle[K, V, C])
- extends BaseShuffleHandle[K, V, C](
- rssShuffleHandleInfo.getShuffleId,
- rssShuffleHandleInfo.getDependency) {}
+case class NativeFilterExec(condition: Expression, override val child: SparkPlan)
+ extends NativeFilterBase(condition, child) {
+
+ override def withNewChildren(newChildren: Seq[SparkPlan]): SparkPlan =
+ copy(child = newChildren.head)
+}
diff --git a/spark-extension-shims-spark241kwaiae/src/main/scala/org/apache/spark/sql/execution/blaze/plan/NativeGenerateExec.scala b/spark-extension-shims-spark241kwaiae/src/main/scala/org/apache/spark/sql/execution/blaze/plan/NativeGenerateExec.scala
new file mode 100644
index 00000000..40ab1db5
--- /dev/null
+++ b/spark-extension-shims-spark241kwaiae/src/main/scala/org/apache/spark/sql/execution/blaze/plan/NativeGenerateExec.scala
@@ -0,0 +1,32 @@
+/*
+ * Copyright 2022 The Blaze Authors
+ *
+ * Licensed 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.execution.blaze.plan
+
+import org.apache.spark.sql.catalyst.expressions.Attribute
+import org.apache.spark.sql.catalyst.expressions.Generator
+import org.apache.spark.sql.execution.SparkPlan
+
+case class NativeGenerateExec(
+ generator: Generator,
+ requiredChildOutput: Seq[Attribute],
+ outer: Boolean,
+ generatorOutput: Seq[Attribute],
+ override val child: SparkPlan)
+ extends NativeGenerateBase(generator, requiredChildOutput, outer, generatorOutput, child) {
+
+ override def withNewChildren(newChildren: Seq[SparkPlan]): SparkPlan =
+ copy(child = newChildren.head)
+}
diff --git a/spark-extension-shims-spark3/src/main/scala/org/apache/spark/shuffle/uniffle/RssShuffleHandleWrapper.scala b/spark-extension-shims-spark241kwaiae/src/main/scala/org/apache/spark/sql/execution/blaze/plan/NativeGlobalLimitExec.scala
similarity index 64%
copy from spark-extension-shims-spark3/src/main/scala/org/apache/spark/shuffle/uniffle/RssShuffleHandleWrapper.scala
copy to spark-extension-shims-spark241kwaiae/src/main/scala/org/apache/spark/sql/execution/blaze/plan/NativeGlobalLimitExec.scala
index e8486b0a..429d93bb 100644
--- a/spark-extension-shims-spark3/src/main/scala/org/apache/spark/shuffle/uniffle/RssShuffleHandleWrapper.scala
+++ b/spark-extension-shims-spark241kwaiae/src/main/scala/org/apache/spark/sql/execution/blaze/plan/NativeGlobalLimitExec.scala
@@ -13,11 +13,13 @@
* See the License for the specific language governing permissions and
* limitations under the License.
*/
-package org.apache.spark.shuffle.uniffle
+package org.apache.spark.sql.execution.blaze.plan
-import org.apache.spark.shuffle.{BaseShuffleHandle, RssShuffleHandle}
+import org.apache.spark.sql.execution.SparkPlan
-class RssShuffleHandleWrapper[K, V, C](val rssShuffleHandleInfo: RssShuffleHandle[K, V, C])
- extends BaseShuffleHandle[K, V, C](
- rssShuffleHandleInfo.getShuffleId,
- rssShuffleHandleInfo.getDependency) {}
+case class NativeGlobalLimitExec(limit: Long, override val child: SparkPlan)
+ extends NativeGlobalLimitBase(limit, child) {
+
+ override def withNewChildren(newChildren: Seq[SparkPlan]): SparkPlan =
+ copy(child = newChildren.head)
+}
diff --git a/spark-extension-shims-spark3/src/main/scala/org/apache/spark/shuffle/uniffle/RssShuffleHandleWrapper.scala b/spark-extension-shims-spark241kwaiae/src/main/scala/org/apache/spark/sql/execution/blaze/plan/NativeLocalLimitExec.scala
similarity index 64%
copy from spark-extension-shims-spark3/src/main/scala/org/apache/spark/shuffle/uniffle/RssShuffleHandleWrapper.scala
copy to spark-extension-shims-spark241kwaiae/src/main/scala/org/apache/spark/sql/execution/blaze/plan/NativeLocalLimitExec.scala
index e8486b0a..094871c4 100644
--- a/spark-extension-shims-spark3/src/main/scala/org/apache/spark/shuffle/uniffle/RssShuffleHandleWrapper.scala
+++ b/spark-extension-shims-spark241kwaiae/src/main/scala/org/apache/spark/sql/execution/blaze/plan/NativeLocalLimitExec.scala
@@ -13,11 +13,13 @@
* See the License for the specific language governing permissions and
* limitations under the License.
*/
-package org.apache.spark.shuffle.uniffle
+package org.apache.spark.sql.execution.blaze.plan
-import org.apache.spark.shuffle.{BaseShuffleHandle, RssShuffleHandle}
+import org.apache.spark.sql.execution.SparkPlan
-class RssShuffleHandleWrapper[K, V, C](val rssShuffleHandleInfo: RssShuffleHandle[K, V, C])
- extends BaseShuffleHandle[K, V, C](
- rssShuffleHandleInfo.getShuffleId,
- rssShuffleHandleInfo.getDependency) {}
+case class NativeLocalLimitExec(limit: Long, override val child: SparkPlan)
+ extends NativeLocalLimitBase(limit, child) {
+
+ override def withNewChildren(newChildren: Seq[SparkPlan]): SparkPlan =
+ copy(child = newChildren.head)
+}
diff --git a/spark-extension-shims-spark3/src/main/scala/org/apache/spark/shuffle/uniffle/RssShuffleHandleWrapper.scala b/spark-extension-shims-spark241kwaiae/src/main/scala/org/apache/spark/sql/execution/blaze/plan/NativeOrcScanExec.scala
similarity index 57%
copy from spark-extension-shims-spark3/src/main/scala/org/apache/spark/shuffle/uniffle/RssShuffleHandleWrapper.scala
copy to spark-extension-shims-spark241kwaiae/src/main/scala/org/apache/spark/sql/execution/blaze/plan/NativeOrcScanExec.scala
index e8486b0a..1c1120c5 100644
--- a/spark-extension-shims-spark3/src/main/scala/org/apache/spark/shuffle/uniffle/RssShuffleHandleWrapper.scala
+++ b/spark-extension-shims-spark241kwaiae/src/main/scala/org/apache/spark/sql/execution/blaze/plan/NativeOrcScanExec.scala
@@ -13,11 +13,15 @@
* See the License for the specific language governing permissions and
* limitations under the License.
*/
-package org.apache.spark.shuffle.uniffle
+package org.apache.spark.sql.execution.blaze.plan
-import org.apache.spark.shuffle.{BaseShuffleHandle, RssShuffleHandle}
+import org.apache.spark.sql.execution.FileSourceScanExec
+import org.apache.spark.sql.execution.statsEstimation.Statistics
-class RssShuffleHandleWrapper[K, V, C](val rssShuffleHandleInfo: RssShuffleHandle[K, V, C])
- extends BaseShuffleHandle[K, V, C](
- rssShuffleHandleInfo.getShuffleId,
- rssShuffleHandleInfo.getDependency) {}
+case class NativeOrcScanExec(basedFileScan: FileSourceScanExec)
+ extends NativeOrcScanBase(basedFileScan) {
+ override def computeStats(): Statistics = basedFileScan.computeStats()
+
+ override def simpleString: String =
+ s"$nodeName (${basedFileScan.simpleString})"
+}
diff --git a/spark-extension-shims-spark241kwaiae/src/main/scala/org/apache/spark/sql/execution/blaze/plan/NativeParquetInsertIntoHiveTableExec.scala b/spark-extension-shims-spark241kwaiae/src/main/scala/org/apache/spark/sql/execution/blaze/plan/NativeParquetInsertIntoHiveTableExec.scala
new file mode 100644
index 00000000..d31f65e4
--- /dev/null
+++ b/spark-extension-shims-spark241kwaiae/src/main/scala/org/apache/spark/sql/execution/blaze/plan/NativeParquetInsertIntoHiveTableExec.scala
@@ -0,0 +1,151 @@
+/*
+ * Copyright 2022 The Blaze Authors
+ *
+ * Licensed 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.execution.blaze.plan
+
+import org.apache.commons.lang3.reflect.FieldUtils
+import org.apache.hadoop.conf.Configuration
+import org.apache.spark.sql.blaze.Shims
+import org.apache.spark.sql.execution.SparkPlan
+import org.apache.spark.sql.hive.execution.InsertIntoHiveTable
+import org.apache.spark.sql.SparkSession
+import org.apache.spark.sql.execution.metric.SQLMetric
+import org.apache.spark.sql.Row
+import org.apache.spark.sql.catalyst.InternalRow
+import org.apache.spark.sql.catalyst.catalog.CatalogTable
+import org.apache.spark.sql.catalyst.expressions.Attribute
+import org.apache.spark.sql.catalyst.plans.logical.LogicalPlan
+import org.apache.spark.sql.execution.datasources.{BasicWriteJobStatsTracker, PartitionWriteJobStatsTracker, PartitionWriteTaskStatsTracker, WriteTaskStatsTracker}
+import org.apache.spark.sql.types.StructType
+import org.apache.spark.util.SerializableConfiguration
+
+case class NativeParquetInsertIntoHiveTableExec(
+ cmd: InsertIntoHiveTable,
+ override val child: SparkPlan)
+ extends NativeParquetInsertIntoHiveTableBase(cmd, child) {
+
+ override protected def getInsertIntoHiveTableCommand(
+ table: CatalogTable,
+ partition: Map[String, Option[String]],
+ query: LogicalPlan,
+ overwrite: Boolean,
+ ifPartitionNotExists: Boolean,
+ outputColumnNames: Seq[String],
+ metrics: Map[String, SQLMetric]): InsertIntoHiveTable = {
+ new BlazeInsertIntoHiveTable30(
+ table,
+ partition,
+ query,
+ overwrite,
+ ifPartitionNotExists,
+ outputColumnNames,
+ metrics)
+ }
+
+ override def withNewChildren(newChildren: Seq[SparkPlan]): SparkPlan =
+ copy(child = newChildren.head)
+
+ class BlazeInsertIntoHiveTable30(
+ table: CatalogTable,
+ partition: Map[String, Option[String]],
+ query: LogicalPlan,
+ overwrite: Boolean,
+ ifPartitionNotExists: Boolean,
+ outputColumnNames: Seq[String],
+ outerMetrics: Map[String, SQLMetric])
+ extends InsertIntoHiveTable(
+ table,
+ partition,
+ query,
+ overwrite,
+ ifPartitionNotExists,
+ outputColumnNames) {
+
+ override lazy val metrics: Map[String, SQLMetric] = outerMetrics
+
+ override def run(sparkSession: SparkSession, child: SparkPlan): Seq[Row] = {
+ val nativeParquetSink =
+ Shims.get.createNativeParquetSinkExec(sparkSession, table, partition, child, metrics)
+ super.run(sparkSession, nativeParquetSink)
+ }
+
+ override def basicWriteJobStatsTracker(hadoopConf: org.apache.hadoop.conf.Configuration) = {
+ import org.apache.spark.sql.catalyst.InternalRow
+ import org.apache.spark.sql.execution.datasources.BasicWriteJobStatsTracker
+ import org.apache.spark.sql.execution.datasources.BasicWriteTaskStatsTracker
+ import org.apache.spark.sql.execution.datasources.WriteTaskStatsTracker
+ import org.apache.spark.util.SerializableConfiguration
+
+ val serializableHadoopConf = new SerializableConfiguration(hadoopConf)
+ new BasicWriteJobStatsTracker(serializableHadoopConf, metrics) {
+ override def newTaskInstance(): WriteTaskStatsTracker = {
+ new BasicWriteTaskStatsTracker(serializableHadoopConf.value) {
+ override def newRow(row: InternalRow): Unit = {
+ if (!ParquetSinkTaskContext.get.isNative) {
+ return super.newRow(row)
+ }
+ }
+
+ override def statCurrentFile(isLast: Boolean): Unit = {
+ if (!ParquetSinkTaskContext.get.isNative) {
+ return super.statCurrentFile(isLast)
+ }
+
+ curFile.foreach { _ =>
+ val outputFileStat = ParquetSinkTaskContext.get.processedOutputFiles.remove()
+ this.numRows += outputFileStat.numRows
+ }
+ super.statCurrentFile(isLast)
+ }
+ }
+ }
+ }
+ }
+
+ override def partitionWriteJobStatsTracker(
+ hadoopConf: Configuration,
+ partitionAttributes: Seq[Attribute],
+ staticPartition: Map[String, Option[String]],
+ partitionSchema: StructType): BasicWriteJobStatsTracker = {
+
+ val serializableHadoopConf = new SerializableConfiguration(hadoopConf)
+ new PartitionWriteJobStatsTracker(serializableHadoopConf, partitionAttributes, metrics, staticPartition, partitionSchema) {
+
+ override def newTaskInstance(): WriteTaskStatsTracker = {
+ new PartitionWriteTaskStatsTracker(partitionAttributes, serializableHadoopConf.value, staticPartition, partitionSchema) {
+ override def newRow(row: InternalRow): Unit = {
+ if (!ParquetSinkTaskContext.get.isNative) {
+ super.newRow(row)
+ }
+ }
+
+ override def statCurrentFile(isLast: Boolean): Unit = {
+ if (!ParquetSinkTaskContext.get.isNative) {
+ return super.statCurrentFile(isLast)
+ }
+
+ curFile.foreach { _ =>
+ val outputFileStat = ParquetSinkTaskContext.get.processedOutputFiles.remove()
+ this.numRows += outputFileStat.numRows
+ FieldUtils.writeField(this, "currentPartitionRows", outputFileStat.numRows, true)
+ }
+ super.statCurrentFile(isLast)
+ }
+ }
+ }
+ }
+ }
+ }
+}
diff --git a/spark-extension-shims-spark3/src/main/scala/org/apache/spark/shuffle/uniffle/RssShuffleHandleWrapper.scala b/spark-extension-shims-spark241kwaiae/src/main/scala/org/apache/spark/sql/execution/blaze/plan/NativeParquetScanExec.scala
similarity index 57%
copy from spark-extension-shims-spark3/src/main/scala/org/apache/spark/shuffle/uniffle/RssShuffleHandleWrapper.scala
copy to spark-extension-shims-spark241kwaiae/src/main/scala/org/apache/spark/sql/execution/blaze/plan/NativeParquetScanExec.scala
index e8486b0a..b4feb7dc 100644
--- a/spark-extension-shims-spark3/src/main/scala/org/apache/spark/shuffle/uniffle/RssShuffleHandleWrapper.scala
+++ b/spark-extension-shims-spark241kwaiae/src/main/scala/org/apache/spark/sql/execution/blaze/plan/NativeParquetScanExec.scala
@@ -13,11 +13,15 @@
* See the License for the specific language governing permissions and
* limitations under the License.
*/
-package org.apache.spark.shuffle.uniffle
+package org.apache.spark.sql.execution.blaze.plan
-import org.apache.spark.shuffle.{BaseShuffleHandle, RssShuffleHandle}
+import org.apache.spark.sql.execution.FileSourceScanExec
+import org.apache.spark.sql.execution.statsEstimation.Statistics
-class RssShuffleHandleWrapper[K, V, C](val rssShuffleHandleInfo: RssShuffleHandle[K, V, C])
- extends BaseShuffleHandle[K, V, C](
- rssShuffleHandleInfo.getShuffleId,
- rssShuffleHandleInfo.getDependency) {}
+case class NativeParquetScanExec(basedFileScan: FileSourceScanExec)
+ extends NativeParquetScanBase(basedFileScan) {
+ override def computeStats(): Statistics = basedFileScan.computeStats()
+
+ override def simpleString: String =
+ s"$nodeName (${basedFileScan.simpleString})"
+}
diff --git a/spark-extension-shims-spark241kwaiae/src/main/scala/org/apache/spark/sql/execution/blaze/plan/NativeParquetSinkExec.scala b/spark-extension-shims-spark241kwaiae/src/main/scala/org/apache/spark/sql/execution/blaze/plan/NativeParquetSinkExec.scala
new file mode 100644
index 00000000..b61a29ff
--- /dev/null
+++ b/spark-extension-shims-spark241kwaiae/src/main/scala/org/apache/spark/sql/execution/blaze/plan/NativeParquetSinkExec.scala
@@ -0,0 +1,57 @@
+/*
+ * Copyright 2022 The Blaze Authors
+ *
+ * Licensed 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.execution.blaze.plan
+
+import org.apache.hadoop.conf.Configuration
+import org.apache.hadoop.hive.ql.plan.TableDesc
+import org.apache.hadoop.io.compress.ZStandardCodec
+import org.apache.parquet.hadoop.metadata.CompressionCodecName
+import org.apache.spark.sql.execution.SparkPlan
+import org.apache.spark.sql.SparkSession
+import org.apache.spark.sql.catalyst.catalog.CatalogTable
+import org.apache.spark.sql.execution.metric.SQLMetric
+import org.apache.spark.sql.hive.execution.HiveOptions
+
+case class NativeParquetSinkExec(
+ sparkSession: SparkSession,
+ table: CatalogTable,
+ partition: Map[String, Option[String]],
+ override val child: SparkPlan,
+ override val metrics: Map[String, SQLMetric])
+ extends NativeParquetSinkBase(sparkSession, table, partition, child, metrics) {
+
+ override def withNewChildren(newChildren: Seq[SparkPlan]): SparkPlan =
+ copy(child = newChildren.head)
+
+ override protected def newHadoopConf(tableDesc: TableDesc): Configuration = {
+ val hadoopConf = super.newHadoopConf(tableDesc)
+
+ // Set compression by priority
+ HiveOptions
+ .getHiveWriteCompression(tableDesc, sparkSession.sessionState.conf)
+ .foreach { case (compression, codec) =>
+ hadoopConf.set(compression, codec)
+ if (codec == CompressionCodecName.ZSTD.name()) {
+ val defaultLevel = ZStandardCodec.getCompressionLevel(hadoopConf)
+ val hiveFileLevel = hadoopConf
+ .getInt("spark.kwai.hivefile.compression.zstd.level", defaultLevel)
+ .min(10) // do not use zstd level > 10
+ hadoopConf.setInt("parquet.compression.codec.zstd.level", hiveFileLevel)
+ }
+ }
+ hadoopConf
+ }
+}
diff --git a/spark-extension-shims-spark3/src/test/scala/org/apache/spark/sql/blaze/BaseBlazeSQLSuite.scala b/spark-extension-shims-spark241kwaiae/src/main/scala/org/apache/spark/sql/execution/blaze/plan/NativePartialTakeOrderedExec.scala
similarity index 51%
copy from spark-extension-shims-spark3/src/test/scala/org/apache/spark/sql/blaze/BaseBlazeSQLSuite.scala
copy to spark-extension-shims-spark241kwaiae/src/main/scala/org/apache/spark/sql/execution/blaze/plan/NativePartialTakeOrderedExec.scala
index 0ad73598..760fe4f5 100644
--- a/spark-extension-shims-spark3/src/test/scala/org/apache/spark/sql/blaze/BaseBlazeSQLSuite.scala
+++ b/spark-extension-shims-spark241kwaiae/src/main/scala/org/apache/spark/sql/execution/blaze/plan/NativePartialTakeOrderedExec.scala
@@ -13,21 +13,19 @@
* See the License for the specific language governing permissions and
* limitations under the License.
*/
-package org.apache.spark.sql.blaze
+package org.apache.spark.sql.execution.blaze.plan
-import org.apache.spark.SparkConf
-import org.apache.spark.sql.test.SharedSparkSession
+import org.apache.spark.sql.catalyst.expressions.SortOrder
+import org.apache.spark.sql.execution.SparkPlan
+import org.apache.spark.sql.execution.metric.SQLMetric
-trait BaseBlazeSQLSuite extends SharedSparkSession {
-
- override protected def sparkConf: SparkConf = {
- super.sparkConf
- .set("spark.sql.extensions", "org.apache.spark.sql.blaze.BlazeSparkSessionExtension")
- .set(
- "spark.shuffle.manager",
- "org.apache.spark.sql.execution.blaze.shuffle.BlazeShuffleManager")
- .set("spark.memory.offHeap.enabled", "false")
- .set("spark.blaze.enable", "true")
- }
+case class NativePartialTakeOrderedExec(
+ limit: Long,
+ sortOrder: Seq[SortOrder],
+ override val child: SparkPlan,
+ override val metrics: Map[String, SQLMetric])
+ extends NativePartialTakeOrderedBase(limit, sortOrder, child, metrics) {
+ override def withNewChildren(newChildren: Seq[SparkPlan]): SparkPlan =
+ copy(child = newChildren.head)
}
diff --git a/spark-extension-shims-spark3/src/main/scala/org/apache/spark/shuffle/uniffle/RssShuffleHandleWrapper.scala b/spark-extension-shims-spark241kwaiae/src/main/scala/org/apache/spark/sql/execution/blaze/plan/NativeProjectExec.scala
similarity index 58%
copy from spark-extension-shims-spark3/src/main/scala/org/apache/spark/shuffle/uniffle/RssShuffleHandleWrapper.scala
copy to spark-extension-shims-spark241kwaiae/src/main/scala/org/apache/spark/sql/execution/blaze/plan/NativeProjectExec.scala
index e8486b0a..4c3e44d6 100644
--- a/spark-extension-shims-spark3/src/main/scala/org/apache/spark/shuffle/uniffle/RssShuffleHandleWrapper.scala
+++ b/spark-extension-shims-spark241kwaiae/src/main/scala/org/apache/spark/sql/execution/blaze/plan/NativeProjectExec.scala
@@ -13,11 +13,14 @@
* See the License for the specific language governing permissions and
* limitations under the License.
*/
-package org.apache.spark.shuffle.uniffle
+package org.apache.spark.sql.execution.blaze.plan
-import org.apache.spark.shuffle.{BaseShuffleHandle, RssShuffleHandle}
+import org.apache.spark.sql.catalyst.expressions.NamedExpression
+import org.apache.spark.sql.execution.SparkPlan
-class RssShuffleHandleWrapper[K, V, C](val rssShuffleHandleInfo: RssShuffleHandle[K, V, C])
- extends BaseShuffleHandle[K, V, C](
- rssShuffleHandleInfo.getShuffleId,
- rssShuffleHandleInfo.getDependency) {}
+case class NativeProjectExec(projectList: Seq[NamedExpression], override val child: SparkPlan)
+ extends NativeProjectBase(projectList, child) {
+
+ override def withNewChildren(newChildren: Seq[SparkPlan]): SparkPlan =
+ copy(child = newChildren.head)
+}
diff --git a/spark-extension-shims-spark3/src/main/scala/org/apache/spark/shuffle/uniffle/RssShuffleHandleWrapper.scala b/spark-extension-shims-spark241kwaiae/src/main/scala/org/apache/spark/sql/execution/blaze/plan/NativeRenameColumnsExec.scala
similarity index 61%
copy from spark-extension-shims-spark3/src/main/scala/org/apache/spark/shuffle/uniffle/RssShuffleHandleWrapper.scala
copy to spark-extension-shims-spark241kwaiae/src/main/scala/org/apache/spark/sql/execution/blaze/plan/NativeRenameColumnsExec.scala
index e8486b0a..fc9ba631 100644
--- a/spark-extension-shims-spark3/src/main/scala/org/apache/spark/shuffle/uniffle/RssShuffleHandleWrapper.scala
+++ b/spark-extension-shims-spark241kwaiae/src/main/scala/org/apache/spark/sql/execution/blaze/plan/NativeRenameColumnsExec.scala
@@ -13,11 +13,13 @@
* See the License for the specific language governing permissions and
* limitations under the License.
*/
-package org.apache.spark.shuffle.uniffle
+package org.apache.spark.sql.execution.blaze.plan
-import org.apache.spark.shuffle.{BaseShuffleHandle, RssShuffleHandle}
+import org.apache.spark.sql.execution.SparkPlan
-class RssShuffleHandleWrapper[K, V, C](val rssShuffleHandleInfo: RssShuffleHandle[K, V, C])
- extends BaseShuffleHandle[K, V, C](
- rssShuffleHandleInfo.getShuffleId,
- rssShuffleHandleInfo.getDependency) {}
+case class NativeRenameColumnsExec(override val child: SparkPlan, renamedColumnNames: Seq[String])
+ extends NativeRenameColumnsBase(child, renamedColumnNames) {
+
+ override def withNewChildren(newChildren: Seq[SparkPlan]): SparkPlan =
+ copy(child = newChildren.head)
+}
diff --git a/spark-extension-shims-spark241kwaiae/src/main/scala/org/apache/spark/sql/execution/blaze/plan/NativeShuffleExchangeExec.scala b/spark-extension-shims-spark241kwaiae/src/main/scala/org/apache/spark/sql/execution/blaze/plan/NativeShuffleExchangeExec.scala
new file mode 100644
index 00000000..00d858b2
--- /dev/null
+++ b/spark-extension-shims-spark241kwaiae/src/main/scala/org/apache/spark/sql/execution/blaze/plan/NativeShuffleExchangeExec.scala
@@ -0,0 +1,236 @@
+/*
+ * Copyright 2022 The Blaze Authors
+ *
+ * Licensed 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.execution.blaze.plan
+
+import scala.collection.mutable
+import org.apache.spark._
+import org.apache.spark.rdd.MapPartitionsRDD
+import org.apache.spark.rdd.RDD
+import org.apache.spark.scheduler.MapStatus
+import org.apache.spark.shuffle.ShuffleWriteMetricsReporter
+import org.apache.spark.shuffle.ShuffleWriteProcessor
+import org.apache.spark.shuffle.ShuffleWriter
+import org.apache.spark.sql.blaze.NativeHelper
+import org.apache.spark.sql.blaze.NativeRDD
+import org.apache.spark.sql.catalyst.InternalRow
+import org.apache.spark.sql.catalyst.errors.attachTree
+import org.apache.spark.sql.catalyst.plans.physical._
+import org.apache.spark.sql.execution._
+import org.apache.spark.sql.execution.metric.SQLMetric
+import org.apache.spark.sql.execution.metric.SQLMetrics
+import org.apache.spark.sql.execution.metric.SQLShuffleReadMetricsReporter
+import org.apache.spark.sql.execution.metric.SQLShuffleWriteMetricsReporter
+import org.apache.spark.sql.execution.SQLExecution.EXECUTION_ID_KEY
+import org.apache.spark.sql.execution.blaze.shuffle.{BlazeShuffleWriter, BlazeRssShuffleWriter}
+import org.apache.spark.sql.internal.SQLConf
+
+case class NativeShuffleExchangeExec(
+ var newPartitioning: Partitioning,
+ override val child: SparkPlan)
+ extends NativeShuffleExchangeBase(newPartitioning, child) {
+
+ val shuffleDescStr =
+ s"$newPartitioning[${child.output.map(_.dataType.simpleString).mkString(",")}]"
+
+ // NOTE: coordinator can be null after serialization/deserialization,
+ // e.g. it can be null on the Executor side
+ override lazy val writeMetrics: Map[String, SQLMetric] =
+ (mutable.LinkedHashMap[String, SQLMetric]() ++
+ SQLShuffleWriteMetricsReporter.createShuffleWriteMetrics(sparkContext, shuffleDescStr) ++
+ mutable.LinkedHashMap(
+ NativeHelper
+ .getDefaultNativeMetrics(sparkContext)
+ .filterKeys(Set(
+ "stage_id",
+ "mem_spill_count",
+ "mem_spill_size",
+ "mem_spill_iotime",
+ "disk_spill_size",
+ "disk_spill_iotime",
+ "shuffle_write_total_time",
+ "shuffle_read_total_time"))
+ .toSeq: _*)).toMap
+
+ override lazy val readMetrics: Map[String, SQLMetric] =
+ SQLShuffleReadMetricsReporter.createShuffleReadMetrics(sparkContext, shuffleDescStr)
+
+ override lazy val metrics: Map[String, SQLMetric] =
+ (mutable.LinkedHashMap[String, SQLMetric]() ++
+ readMetrics ++
+ writeMetrics ++
+ Map(
+ "dataSize" -> SQLMetrics.createSizeMetric(sparkContext, "data size"),
+ "numPartitions" -> SQLMetrics.createMetric(sparkContext, "number of partitions"))).toMap
+
+ private def preparePostShuffleRDD(
+ shuffleDependency: ShuffleDependency[Int, InternalRow, InternalRow],
+ specifiedPartitionStartIndices: Option[Array[Int]] = None): ShuffledRowRDD = {
+ // If an array of partition start indices is provided, we need to use this array
+ // to create the ShuffledRowRDD. Also, we need to update outputPartitioning to
+ // update the number of post-shuffle partitions.
+ specifiedPartitionStartIndices.foreach { indices =>
+ assert(newPartitioning.isInstanceOf[HashPartitioning])
+ newPartitioning = UnknownPartitioning(indices.length)
+ }
+ new ShuffledRowRDD(shuffleDependency, readMetrics, specifiedPartitionStartIndices)
+ }
+
+ /**
+ * Caches the created ShuffleRowRDD so we can reuse that.
+ */
+ private var cachedShuffleRDD: ShuffledRowRDD = _
+
+ protected override def doExecuteNonNative(): RDD[InternalRow] =
+ attachTree(this, "execute") {
+ // Returns the same ShuffleRowRDD if this plan is used by multiple plans.
+ if (cachedShuffleRDD == null) {
+ cachedShuffleRDD = preparePostShuffleRDD(shuffleDependency)
+ }
+ cachedShuffleRDD
+ }
+
+ private var _mapOutputStatistics: MapOutputStatistics = _
+
+ override def mapOutputStatistics: MapOutputStatistics = _mapOutputStatistics
+
+ override def eagerExecute(): RDD[InternalRow] = {
+ if (cachedShuffleRDD == null) {
+ shuffleDependency.rdd.stageId = this.stageId
+ val allOperatorMetrics = new mutable.HashSet[SQLMetric]()
+ metrics.foreach { case (_: String, v: SQLMetric) =>
+ allOperatorMetrics.add(v)
+ }
+ transformUp { case p: SparkPlan =>
+ p.metrics.foreach { case (_: String, sqlMetric: SQLMetric) =>
+ allOperatorMetrics.add(sqlMetric)
+ }
+ p
+ }
+ val executionId = sqlContext.sparkContext.getLocalProperty(EXECUTION_ID_KEY)
+ SQLMetrics.postDriverMetricUpdates(
+ sqlContext.sparkContext,
+ executionId,
+ allOperatorMetrics.toSeq)
+ if (shuffleDependency.rdd.partitions.length != 0) {
+ // submitMapStage does not accept RDD with 0 partition.
+ // So, we will not submit this dependency.
+ val submittedStageFuture = sqlContext.sparkContext.submitMapStage(shuffleDependency)
+ try {
+ val stat = submittedStageFuture.get()
+ val totalBytesWritten = stat.bytesByPartitionId.sum
+ var dataSizeFactor = 1.0
+
+ val numRecordsWritten = metrics(
+ SQLShuffleWriteMetricsReporter.SHUFFLE_RECORDS_WRITTEN).value
+
+ if (numRecordsWritten > 0 && totalBytesWritten > 0) {
+ val broadcastCountLimit = SQLConf.get.broadcastCountLimit() / 4
+ val broadcastThreshold = SQLConf.get.getConf(SQLConf.AUTO_BROADCASTJOIN_THRESHOLD)
+ val dataSize = metrics("dataSize").value
+
+ if (totalBytesWritten.toDouble <= broadcastThreshold) {
+ if (numRecordsWritten >= broadcastCountLimit) {
+ // NOTE:
+ // in some cases, the number of written records exceeds broadcastCountLimit
+ // but the size is smaller than autoBroadcastThreshold. in this situation
+ // spark incorrectly turns SMJ into BHJ and always fails the broadcast. so we
+ // have to manually increase the stats by setting the dataSizeFactor.
+ dataSizeFactor = broadcastThreshold / totalBytesWritten.toDouble + 0.1
+ }
+ } else {
+ // NOTE:
+ // in some cases, the compression ratio of vectorized data is too high, resulting
+ // the number of tasks in next stage is too small.
+ // increase the dataSizeFactor to avoid this.
+ val minCompressionRatio = sparkContext.conf
+ .getDouble("spark.blaze.shuffle.minCompressRatio", defaultValue = 0.1)
+ val realCompressionRatio = totalBytesWritten.toDouble / dataSize
+ if (realCompressionRatio < minCompressionRatio) {
+ dataSizeFactor = minCompressionRatio / realCompressionRatio
+ }
+ }
+ }
+
+ logInfo(
+ s"shuffleId=${shuffleDependency.shuffleId}" +
+ s", numRecordsWritten=$numRecordsWritten" +
+ s", totalBytesWritten=$totalBytesWritten")
+ _mapOutputStatistics = new MapOutputStatistics(
+ stat.shuffleId,
+ stat.bytesByPartitionId.map(n => (n * dataSizeFactor).ceil.toLong),
+ stat.recordsByPartitionId)
+
+ } catch {
+ case throwable: Throwable =>
+ sparkContext.dagScheduler.sendFailedReason(throwable)
+ throw throwable
+ }
+ }
+ cachedShuffleRDD = preparePostShuffleRDD(shuffleDependency)
+ }
+ cachedShuffleRDD
+ }
+
+ override def createNativeShuffleWriteProcessor(
+ metrics: Map[String, SQLMetric],
+ numPartitions: Int): ShuffleWriteProcessor = {
+
+ new ShuffleWriteProcessor {
+ override protected def createMetricsReporter(
+ context: TaskContext): ShuffleWriteMetricsReporter = {
+ new SQLShuffleWriteMetricsReporter(context.taskMetrics().shuffleWriteMetrics, metrics)
+ }
+
+ override def write(
+ rdd: RDD[_],
+ dep: ShuffleDependency[_, _, _],
+ mapId: Int,
+ context: TaskContext,
+ partition: Partition): MapStatus = {
+
+ val manager = SparkEnv.get.shuffleManager
+ var writer: ShuffleWriter[Any, Any] = null
+ writer = manager.getWriter[Any, Any](
+ dep.shuffleHandle,
+ mapId,
+ context,
+ createMetricsReporter(context))
+ if (SparkEnv.get.conf.get("spark.shuffle.manager", "sort").equals("unify")) {
+ writer
+ .asInstanceOf[BlazeRssShuffleWriter[_, _]]
+ .nativeRssShuffleWrite(
+ rdd.asInstanceOf[MapPartitionsRDD[_, _]].prev.asInstanceOf[NativeRDD],
+ dep,
+ mapId,
+ context,
+ partition,
+ numPartitions)
+
+ } else {
+ writer
+ .asInstanceOf[BlazeShuffleWriter[_, _]]
+ .nativeShuffleWrite(
+ rdd.asInstanceOf[MapPartitionsRDD[_, _]].prev.asInstanceOf[NativeRDD],
+ dep,
+ mapId,
+ context,
+ partition)
+ }
+ writer.stop(true).get
+ }
+ }
+ }
+}
diff --git a/spark-extension-shims-spark3/src/main/scala/org/apache/spark/shuffle/uniffle/RssShuffleHandleWrapper.scala b/spark-extension-shims-spark241kwaiae/src/main/scala/org/apache/spark/sql/execution/blaze/plan/NativeSortExec.scala
similarity index 58%
copy from spark-extension-shims-spark3/src/main/scala/org/apache/spark/shuffle/uniffle/RssShuffleHandleWrapper.scala
copy to spark-extension-shims-spark241kwaiae/src/main/scala/org/apache/spark/sql/execution/blaze/plan/NativeSortExec.scala
index e8486b0a..fbd7e354 100644
--- a/spark-extension-shims-spark3/src/main/scala/org/apache/spark/shuffle/uniffle/RssShuffleHandleWrapper.scala
+++ b/spark-extension-shims-spark241kwaiae/src/main/scala/org/apache/spark/sql/execution/blaze/plan/NativeSortExec.scala
@@ -13,11 +13,16 @@
* See the License for the specific language governing permissions and
* limitations under the License.
*/
-package org.apache.spark.shuffle.uniffle
+package org.apache.spark.sql.execution.blaze.plan
-import org.apache.spark.shuffle.{BaseShuffleHandle, RssShuffleHandle}
+import org.apache.spark.sql.catalyst.expressions.SortOrder
+import org.apache.spark.sql.execution.SparkPlan
-class RssShuffleHandleWrapper[K, V, C](val rssShuffleHandleInfo: RssShuffleHandle[K, V, C])
- extends BaseShuffleHandle[K, V, C](
- rssShuffleHandleInfo.getShuffleId,
- rssShuffleHandleInfo.getDependency) {}
+case class NativeSortExec(
+ sortOrder: Seq[SortOrder],
+ global: Boolean,
+ override val child: SparkPlan)
+ extends NativeSortBase(sortOrder, global, child) {
+ override def withNewChildren(newChildren: Seq[SparkPlan]): SparkPlan =
+ copy(child = newChildren.head)
+}
diff --git a/spark-extension-shims-spark241kwaiae/src/main/scala/org/apache/spark/sql/execution/blaze/plan/NativeSortMergeJoinExec.scala b/spark-extension-shims-spark241kwaiae/src/main/scala/org/apache/spark/sql/execution/blaze/plan/NativeSortMergeJoinExec.scala
new file mode 100644
index 00000000..baee7226
--- /dev/null
+++ b/spark-extension-shims-spark241kwaiae/src/main/scala/org/apache/spark/sql/execution/blaze/plan/NativeSortMergeJoinExec.scala
@@ -0,0 +1,39 @@
+/*
+ * Copyright 2022 The Blaze Authors
+ *
+ * Licensed 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.execution.blaze.plan
+
+import org.apache.spark.sql.catalyst.expressions.Expression
+import org.apache.spark.sql.catalyst.plans.JoinType
+import org.apache.spark.sql.execution.SparkPlan
+import org.apache.spark.sql.execution.joins.SortMergeJoinExec
+
+case class NativeSortMergeJoinExec(
+ override val left: SparkPlan,
+ override val right: SparkPlan,
+ leftKeys: Seq[Expression],
+ rightKeys: Seq[Expression],
+ joinType: JoinType)
+ extends {
+ override val (output, outputPartitioning, outputOrdering) = {
+ val smj = SortMergeJoinExec(leftKeys, rightKeys, joinType, None, left, right)
+ (smj.output, smj.outputPartitioning, smj.outputOrdering)
+ }
+ }
+ with NativeSortMergeJoinBase(left, right, leftKeys, rightKeys, joinType) {
+
+ override def withNewChildren(newChildren: Seq[SparkPlan]): SparkPlan =
+ copy(left = newChildren(0), right = newChildren(1))
+}
diff --git a/spark-extension-shims-spark3/src/main/scala/org/apache/spark/shuffle/uniffle/RssShuffleHandleWrapper.scala b/spark-extension-shims-spark241kwaiae/src/main/scala/org/apache/spark/sql/execution/blaze/plan/NativeTakeOrderedExec.scala
similarity index 57%
copy from spark-extension-shims-spark3/src/main/scala/org/apache/spark/shuffle/uniffle/RssShuffleHandleWrapper.scala
copy to spark-extension-shims-spark241kwaiae/src/main/scala/org/apache/spark/sql/execution/blaze/plan/NativeTakeOrderedExec.scala
index e8486b0a..935462e3 100644
--- a/spark-extension-shims-spark3/src/main/scala/org/apache/spark/shuffle/uniffle/RssShuffleHandleWrapper.scala
+++ b/spark-extension-shims-spark241kwaiae/src/main/scala/org/apache/spark/sql/execution/blaze/plan/NativeTakeOrderedExec.scala
@@ -13,11 +13,17 @@
* See the License for the specific language governing permissions and
* limitations under the License.
*/
-package org.apache.spark.shuffle.uniffle
+package org.apache.spark.sql.execution.blaze.plan
-import org.apache.spark.shuffle.{BaseShuffleHandle, RssShuffleHandle}
+import org.apache.spark.sql.catalyst.expressions.SortOrder
+import org.apache.spark.sql.execution.SparkPlan
-class RssShuffleHandleWrapper[K, V, C](val rssShuffleHandleInfo: RssShuffleHandle[K, V, C])
- extends BaseShuffleHandle[K, V, C](
- rssShuffleHandleInfo.getShuffleId,
- rssShuffleHandleInfo.getDependency) {}
+case class NativeTakeOrderedExec(
+ limit: Long,
+ sortOrder: Seq[SortOrder],
+ override val child: SparkPlan)
+ extends NativeTakeOrderedBase(limit, sortOrder, child) {
+
+ override def withNewChildren(newChildren: Seq[SparkPlan]): SparkPlan =
+ copy(child = newChildren.head)
+}
diff --git a/spark-extension-shims-spark3/src/main/scala/org/apache/spark/shuffle/uniffle/RssShuffleHandleWrapper.scala b/spark-extension-shims-spark241kwaiae/src/main/scala/org/apache/spark/sql/execution/blaze/plan/NativeUnionExec.scala
similarity index 58%
copy from spark-extension-shims-spark3/src/main/scala/org/apache/spark/shuffle/uniffle/RssShuffleHandleWrapper.scala
copy to spark-extension-shims-spark241kwaiae/src/main/scala/org/apache/spark/sql/execution/blaze/plan/NativeUnionExec.scala
index e8486b0a..3c32846a 100644
--- a/spark-extension-shims-spark3/src/main/scala/org/apache/spark/shuffle/uniffle/RssShuffleHandleWrapper.scala
+++ b/spark-extension-shims-spark241kwaiae/src/main/scala/org/apache/spark/sql/execution/blaze/plan/NativeUnionExec.scala
@@ -13,11 +13,16 @@
* See the License for the specific language governing permissions and
* limitations under the License.
*/
-package org.apache.spark.shuffle.uniffle
+package org.apache.spark.sql.execution.blaze.plan
-import org.apache.spark.shuffle.{BaseShuffleHandle, RssShuffleHandle}
+import org.apache.spark.sql.catalyst.expressions.Attribute
+import org.apache.spark.sql.execution.SparkPlan
-class RssShuffleHandleWrapper[K, V, C](val rssShuffleHandleInfo: RssShuffleHandle[K, V, C])
- extends BaseShuffleHandle[K, V, C](
- rssShuffleHandleInfo.getShuffleId,
- rssShuffleHandleInfo.getDependency) {}
+case class NativeUnionExec(
+ override val children: Seq[SparkPlan],
+ override val output: Seq[Attribute])
+ extends NativeUnionBase(children, output) {
+
+ override def withNewChildren(newChildren: Seq[SparkPlan]): SparkPlan =
+ copy(children = newChildren)
+}
diff --git a/spark-extension-shims-spark241kwaiae/src/main/scala/org/apache/spark/sql/execution/blaze/plan/NativeWindowExec.scala b/spark-extension-shims-spark241kwaiae/src/main/scala/org/apache/spark/sql/execution/blaze/plan/NativeWindowExec.scala
new file mode 100644
index 00000000..7537e54c
--- /dev/null
+++ b/spark-extension-shims-spark241kwaiae/src/main/scala/org/apache/spark/sql/execution/blaze/plan/NativeWindowExec.scala
@@ -0,0 +1,32 @@
+/*
+ * Copyright 2022 The Blaze Authors
+ *
+ * Licensed 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.execution.blaze.plan
+
+import org.apache.spark.sql.catalyst.expressions.Expression
+import org.apache.spark.sql.catalyst.expressions.NamedExpression
+import org.apache.spark.sql.catalyst.expressions.SortOrder
+import org.apache.spark.sql.execution.SparkPlan
+
+case class NativeWindowExec(
+ windowExpression: Seq[NamedExpression],
+ partitionSpec: Seq[Expression],
+ orderSpec: Seq[SortOrder],
+ override val child: SparkPlan)
+ extends NativeWindowBase(windowExpression, partitionSpec, orderSpec, None, child) {
+
+ override def withNewChildren(newChildren: Seq[SparkPlan]): SparkPlan =
+ copy(child = newChildren.head)
+}
diff --git a/spark-extension-shims-spark241kwaiae/src/main/scala/org/apache/spark/sql/execution/blaze/shuffle/BlazeBlockStoreShuffleReader.scala b/spark-extension-shims-spark241kwaiae/src/main/scala/org/apache/spark/sql/execution/blaze/shuffle/BlazeBlockStoreShuffleReader.scala
new file mode 100644
index 00000000..bac1d62b
--- /dev/null
+++ b/spark-extension-shims-spark241kwaiae/src/main/scala/org/apache/spark/sql/execution/blaze/shuffle/BlazeBlockStoreShuffleReader.scala
@@ -0,0 +1,87 @@
+/*
+ * Copyright 2022 The Blaze Authors
+ *
+ * Licensed 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.execution.blaze.shuffle
+
+import java.io.InputStream
+
+import org.apache.spark.MapOutputTracker
+import org.apache.spark.SparkEnv
+import org.apache.spark.TaskContext
+
+import org.apache.spark.internal.Logging
+import org.apache.spark.internal.config
+import org.apache.spark.shuffle.BaseShuffleHandle
+import org.apache.spark.shuffle.ShuffleReadMetricsReporter
+import org.apache.spark.storage.BlockId
+import org.apache.spark.storage.BlockManager
+import org.apache.spark.storage.ShuffleBlockFetcherIterator
+
+class BlazeBlockStoreShuffleReader[K, C](
+ handle: BaseShuffleHandle[K, _, C],
+ startPartition: Int,
+ endPartition: Int,
+ context: TaskContext,
+ readMetrics: ShuffleReadMetricsReporter,
+ blockManager: BlockManager = SparkEnv.get.blockManager,
+ mapOutputTracker: MapOutputTracker = SparkEnv.get.mapOutputTracker,
+ startMapId: Option[Int] = None,
+ endMapId: Option[Int] = None,
+ isLocalShuffleReader: Boolean = false)
+ extends BlazeBlockStoreShuffleReaderBase[K, C](handle, context)
+ with Logging {
+
+ override def readBlocks(): Iterator[InputStream] = {
+ val start = System.currentTimeMillis()
+ val blocksByAddress = (startMapId, endMapId) match {
+ case (Some(startId), Some(endId)) =>
+ mapOutputTracker.getMapSizesByExecutorId(
+ handle.shuffleId,
+ startPartition,
+ endPartition,
+ startId,
+ endId,
+ dep.serializer.supportsRelocationOfSerializedObjects)
+ case (None, None) =>
+ mapOutputTracker.getMapSizesByExecutorId(
+ handle.shuffleId,
+ startPartition,
+ endPartition,
+ dep.serializer.supportsRelocationOfSerializedObjects)
+ case (_, _) =>
+ throw new IllegalArgumentException("startMapId and endMapId should be both set or unset")
+ }
+ val end = System.currentTimeMillis()
+ context.taskMetrics().executionPhaseMetric.incProduceFetchRequestTime(end - start)
+
+ val blockStreams = new ShuffleBlockFetcherIterator(
+ context,
+ blockManager.shuffleClient,
+ blockManager,
+ blocksByAddress,
+ streamWrapper = (_, in) => in,
+ // Note: we use getSizeAsMb when no suffix is provided for backwards compatibility
+ SparkEnv.get.conf.getSizeAsMb("spark.reducer.maxSizeInFlight", "48m") * 1024 * 1024,
+ SparkEnv.get.conf.getInt("spark.reducer.maxReqsInFlight", Int.MaxValue),
+ SparkEnv.get.conf.get(config.REDUCER_MAX_BLOCKS_IN_FLIGHT_PER_ADDRESS),
+ SparkEnv.get.conf.get(config.MAX_REMOTE_BLOCK_SIZE_FETCH_TO_MEM),
+ SparkEnv.get.conf.getBoolean("spark.shuffle.detectCorrupt", true),
+ readMetrics,
+ localShuffle = isLocalShuffleReader &&
+ SparkEnv.get.conf.getBoolean("spark.kwai.localShuffle.readHdfs.enabled", false))
+
+ blockStreams.map(_._2)
+ }
+}
diff --git a/spark-extension-shims-spark241kwaiae/src/main/scala/org/apache/spark/sql/execution/blaze/shuffle/BlazeRssBlockStoreShuffleReader.scala b/spark-extension-shims-spark241kwaiae/src/main/scala/org/apache/spark/sql/execution/blaze/shuffle/BlazeRssBlockStoreShuffleReader.scala
new file mode 100644
index 00000000..ea1fea05
--- /dev/null
+++ b/spark-extension-shims-spark241kwaiae/src/main/scala/org/apache/spark/sql/execution/blaze/shuffle/BlazeRssBlockStoreShuffleReader.scala
@@ -0,0 +1,154 @@
+/*
+ * Copyright 2022 The Blaze Authors
+ *
+ * Licensed 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.execution.blaze.shuffle
+
+import java.io.InputStream
+
+import scala.collection.JavaConverters._
+
+import com.kuaishou.dataarch.shuffle.common.io.UnifiedDataInputBuffer
+import com.kuaishou.dataarch.shuffle.common.iterator.UnifiedRawKeyValueIterator
+import com.kuaishou.dataarch.shuffle.common.reporter.Reporter
+import com.kuaishou.dataarch.shuffle.proto.dto.common.PartitionStatistics
+import com.kuaishou.dataarch.shuffle.proto.dto.enums.AppTypeEnum
+import com.kuaishou.dataarch.shuffle.proto.dto.enums.ShuffleTypeEnum
+import com.kuaishou.dataarch.shuffle.reader.model
+import com.kuaishou.dataarch.shuffle.reader.service.UnifiedShuffleReader
+import com.kuaishou.dataarch.shuffle.reader.service.UnifiedShuffleReader.Context
+import org.apache.avro.util.ReusableByteArrayInputStream
+import org.apache.spark.MapOutputTracker
+import org.apache.spark.SparkEnv
+import org.apache.spark.TaskContext
+
+import org.apache.spark.internal.Logging
+import org.apache.spark.shuffle.stream.StreamShuffleHandle
+import org.apache.spark.shuffle.ShuffleReadMetricsReporter
+import org.apache.spark.storage.BlockId
+import org.apache.spark.storage.BlockManager
+import org.apache.spark.util.CompletionIterator
+import org.apache.spark.util.Utils
+
+class BlazeRssBlockStoreShuffleReader[K, C](
+ handle: StreamShuffleHandle[K, C],
+ startPartition: Int,
+ endPartition: Int,
+ context: TaskContext,
+ readMetrics: ShuffleReadMetricsReporter,
+ blockManager: BlockManager = SparkEnv.get.blockManager,
+ mapOutputTracker: MapOutputTracker = SparkEnv.get.mapOutputTracker,
+ startMapId: Option[Int] = None,
+ endMapId: Option[Int] = None,
+ partitionStatistics: Option[Seq[PartitionStatistics]] = None,
+ isLocalShuffleReader: Boolean = false)
+ extends BlazeBlockStoreShuffleReaderBase[K, C](handle, context)
+ with Logging {
+
+ override def readBlocks(): Iterator[InputStream] = {
+
+ val sessionInfo = handle.sessionInfo
+
+ def resetStream(
+ buffer: UnifiedDataInputBuffer,
+ reuseStream: ReusableByteArrayInputStream): Unit = {
+ reuseStream.setByteArray(buffer.getData, buffer.getPosition, buffer.getLength)
+ readMetrics.incRemoteBytesRead(buffer.getLength)
+ }
+
+ lazy val localDirs = Utils.getConfiguredLocalDirs(SparkEnv.get.conf).mkString(",")
+
+ val reporter = new ReaderMetricsReporter(readMetrics)
+
+ val readContextBuilder = Context
+ .builder[K, C]()
+ .sessionId(sessionInfo.getSessionId)
+ .sessionAttemptId(sessionInfo.getAttemptId)
+ .stageId(handle.shuffleId)
+ .taskId(context.partitionId())
+ .appShuffleMasterId(sessionInfo.getAppShuffleMasterId)
+ .taskAttemptId(context.attemptNumber())
+ .codec(sessionInfo.getCompressClass)
+ .startPartId(startPartition)
+ .endPartId(endPartition)
+ .startMapId(startMapId.getOrElse(0))
+ .endMapId(endMapId.getOrElse(0))
+ .reporter(reporter)
+ .appType(AppTypeEnum.SPARK)
+ .shuffleType(ShuffleTypeEnum.HASH)
+ .localDirsStr(localDirs)
+ .conf(Utils.getHadoopConf())
+ partitionStatistics.foreach(it => readContextBuilder.partitionStatistics(it.asJava))
+
+ val readContext = readContextBuilder.build()
+
+ val realReader = UnifiedShuffleReader.getReader(readContext)
+
+ var shuffleNextCalled: Long = 0
+
+ val blockIter: Iterator[(BlockId, InputStream)] = new Iterator[(BlockId, InputStream)] {
+ val start: Long = System.currentTimeMillis()
+ val innerIterator: UnifiedRawKeyValueIterator = realReader.read()
+ readMetrics.incBytesReadTime(System.currentTimeMillis() - start)
+ val rbais = new ReusableByteArrayInputStream()
+ var innerHashNext = false
+
+ var nextCalled = false
+ // call inner next in constructor for the first element
+ innerNext()
+
+ def innerNext(): Unit = {
+ val start = System.currentTimeMillis()
+ innerHashNext = innerIterator.next()
+ shuffleNextCalled += 1
+ readMetrics.incFetchWaitTime(System.currentTimeMillis() - start)
+ }
+
+ def getBlock: (BlockId, InputStream) = {
+ resetStream(innerIterator.getValue, rbais)
+ readMetrics.incRemoteBlocksFetched(1)
+ (null, rbais)
+ }
+
+ override def hasNext: Boolean = {
+ if (nextCalled) {
+ innerNext()
+ nextCalled = false
+ }
+ innerHashNext
+ }
+
+ override def next(): (BlockId, InputStream) = {
+ nextCalled = true
+ getBlock
+
+ }
+ }
+
+ val succTaskIter =
+ CompletionIterator[(BlockId, InputStream), Iterator[(BlockId, InputStream)]](
+ blockIter, {
+ logInfo(s"Closing stream shuffle reader in task ${context.taskAttemptId()} for now...")
+ realReader.close()
+ logInfo(s"Stream shuffle reader in task ${context.taskAttemptId()} closed.")
+ context.taskMetrics().mergeShuffleReadMetrics(true)
+ })
+ succTaskIter.map(_._2)
+ }
+
+ class ReaderMetricsReporter(readMetrics: ShuffleReadMetricsReporter)
+ extends Reporter[model.ShuffleStatus] {
+ override def report(t: model.ShuffleStatus): Unit = {}
+ }
+}
diff --git a/spark-extension-shims-spark241kwaiae/src/main/scala/org/apache/spark/sql/execution/blaze/shuffle/BlazeRssShuffleManager.scala b/spark-extension-shims-spark241kwaiae/src/main/scala/org/apache/spark/sql/execution/blaze/shuffle/BlazeRssShuffleManager.scala
new file mode 100644
index 00000000..45843202
--- /dev/null
+++ b/spark-extension-shims-spark241kwaiae/src/main/scala/org/apache/spark/sql/execution/blaze/shuffle/BlazeRssShuffleManager.scala
@@ -0,0 +1,130 @@
+/*
+ * Copyright 2022 The Blaze Authors
+ *
+ * Licensed 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.execution.blaze.shuffle
+
+import java.util.concurrent.ConcurrentHashMap
+
+import com.kuaishou.dataarch.shuffle.proto.dto.common.PartitionStatistics
+import org.apache.spark.SparkConf
+import org.apache.spark.TaskContext
+
+import org.apache.spark.shuffle._
+import org.apache.spark.shuffle.stream.StreamShuffleHandle
+import org.apache.spark.shuffle.stream.StreamShuffleManager
+import org.apache.spark.sql.blaze.kwai.KwaiPrivilegedHDFSBlockManager
+import org.apache.spark.sql.execution.blaze.shuffle.BlazeShuffleDependency.isArrowShuffle
+
+class BlazeRssShuffleManager(conf: SparkConf) extends UnifiedShuffleManager(conf) {
+
+ // introduce a wrapped HDFSBlockManager because we found some permission issues
+ // if the original HDFSBlockManager is initialized inside a native thread.
+ KwaiPrivilegedHDFSBlockManager.setup(conf)
+
+ /**
+ * Get a reader for a range of reduce partitions (startPartition to endPartition-1, inclusive).
+ * Called on executors by reduce tasks.
+ */
+ override def getReader[K, C](
+ handle: ShuffleHandle,
+ startPartition: Int,
+ endPartition: Int,
+ context: TaskContext,
+ metrics: ShuffleReadMetricsReporter): ShuffleReader[K, C] = {
+
+ if (isArrowShuffle(handle)) {
+ new BlazeRssBlockStoreShuffleReader[K, C](
+ handle.asInstanceOf[StreamShuffleHandle[K, C]],
+ startPartition,
+ endPartition,
+ context,
+ metrics)
+ } else {
+ getEffectiveManager(handle).getReader(
+ handle,
+ startPartition,
+ endPartition,
+ context,
+ metrics)
+ }
+ }
+
+ override def getReader[K, C](
+ handle: ShuffleHandle,
+ startPartition: Int,
+ endPartition: Int,
+ context: TaskContext,
+ metrics: ShuffleReadMetricsReporter,
+ startMapId: Int,
+ endMapId: Int): ShuffleReader[K, C] = {
+
+ if (isArrowShuffle(handle)) {
+ new BlazeRssBlockStoreShuffleReader[K, C](
+ handle.asInstanceOf[StreamShuffleHandle[K, C]],
+ startPartition,
+ endPartition,
+ context,
+ metrics,
+ startMapId = Some(startMapId),
+ endMapId = Some(endMapId))
+ } else {
+ getEffectiveManager(handle).getReader(
+ handle,
+ startPartition,
+ endPartition,
+ context,
+ metrics,
+ startMapId,
+ endMapId)
+ }
+ }
+
+ override def getRebalancedReader[K, C](
+ handle: ShuffleHandle,
+ partitionId: Int,
+ context: TaskContext,
+ metrics: ShuffleReadMetricsReporter,
+ partitionStatistics: Seq[PartitionStatistics]): ShuffleReader[K, C] = {
+ if (isArrowShuffle(handle)) {
+ new BlazeRssBlockStoreShuffleReader[K, C](
+ handle.asInstanceOf[StreamShuffleHandle[K, C]],
+ partitionId,
+ partitionId + 1,
+ context,
+ metrics,
+ startMapId = None,
+ endMapId = None,
+ partitionStatistics = Some(partitionStatistics))
+ } else {
+ getEffectiveManager(handle)
+ .asInstanceOf[StreamShuffleManager]
+ .getRebalancedReader(handle, partitionId, context, metrics, partitionStatistics)
+ }
+ }
+
+ /** Get a writer for a given partition. Called on executors by map tasks. */
+ override def getWriter[K, V](
+ handle: ShuffleHandle,
+ mapId: Int,
+ context: TaskContext,
+ metrics: ShuffleWriteMetricsReporter): ShuffleWriter[K, V] = {
+ if (isArrowShuffle(handle)) {
+ new BlazeRssShuffleWriter[K, V](metrics)
+ } else {
+ getEffectiveManager(handle).getWriter(handle, mapId, context, metrics)
+ }
+
+ }
+}
diff --git a/spark-extension-shims-spark241kwaiae/src/main/scala/org/apache/spark/sql/execution/blaze/shuffle/BlazeRssShuffleWriter.scala b/spark-extension-shims-spark241kwaiae/src/main/scala/org/apache/spark/sql/execution/blaze/shuffle/BlazeRssShuffleWriter.scala
new file mode 100644
index 00000000..ff5fde9e
--- /dev/null
+++ b/spark-extension-shims-spark241kwaiae/src/main/scala/org/apache/spark/sql/execution/blaze/shuffle/BlazeRssShuffleWriter.scala
@@ -0,0 +1,49 @@
+/*
+ * Copyright 2022 The Blaze Authors
+ *
+ * Licensed 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.execution.blaze.shuffle
+
+import org.apache.spark.scheduler.MapStatus
+import org.apache.spark.shuffle.ShuffleHandle
+import org.apache.spark.shuffle.ShuffleWriteMetricsReporter
+import org.apache.spark.sql.blaze.Shims
+
+class BlazeRssShuffleWriter[K, V](metrics: ShuffleWriteMetricsReporter)
+ extends BlazeRssShuffleWriterBase[K, V](metrics) {
+
+ var rssShuffleWriterObject: RssPartitionWriterBase = _
+
+ override def getRssPartitionWriter(
+ handle: ShuffleHandle,
+ mapId: Int,
+ metrics: ShuffleWriteMetricsReporter,
+ numPartitions: Int): RssPartitionWriterBase = {
+
+ rssShuffleWriterObject = Shims.get
+ .getRssPartitionWriter(handle, mapId, metrics, numPartitions)
+ .get
+ rssShuffleWriterObject
+ }
+
+ override def rssStop(success: Boolean): Option[MapStatus] = {
+ if (rssShuffleWriterObject != null) {
+ val time = System.currentTimeMillis()
+ rssShuffleWriterObject.close(success)
+ rssShuffleWriterObject = null
+ metrics.incWriteTime(System.currentTimeMillis() - time)
+ }
+ None
+ }
+}
diff --git a/spark-extension-shims-spark241kwaiae/src/main/scala/org/apache/spark/sql/execution/blaze/shuffle/BlazeShuffleManager.scala b/spark-extension-shims-spark241kwaiae/src/main/scala/org/apache/spark/sql/execution/blaze/shuffle/BlazeShuffleManager.scala
new file mode 100644
index 00000000..ec968321
--- /dev/null
+++ b/spark-extension-shims-spark241kwaiae/src/main/scala/org/apache/spark/sql/execution/blaze/shuffle/BlazeShuffleManager.scala
@@ -0,0 +1,127 @@
+/*
+ * Copyright 2022 The Blaze Authors
+ *
+ * Licensed 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.execution.blaze.shuffle
+
+import java.util.concurrent.ConcurrentHashMap
+
+import org.apache.spark.ShuffleDependency
+import org.apache.spark.SparkConf
+import org.apache.spark.TaskContext
+
+import org.apache.spark.internal.Logging
+import org.apache.spark.shuffle._
+import org.apache.spark.shuffle.sort.SortShuffleManager
+import org.apache.spark.sql.blaze.kwai.KwaiPrivilegedHDFSBlockManager
+import org.apache.spark.sql.execution.blaze.shuffle.BlazeShuffleDependency.isArrowShuffle
+
+class BlazeShuffleManager(conf: SparkConf) extends ShuffleManager with Logging {
+
+ // introduce a wrapped HDFSBlockManager because we found some permission issues
+ // if the original HDFSBlockManager is initialized inside a native thread.
+ KwaiPrivilegedHDFSBlockManager.setup(conf)
+
+ val sortShuffleManager = new SortShuffleManager(conf)
+
+ override val shuffleBlockResolver: IndexShuffleBlockResolver =
+ sortShuffleManager.shuffleBlockResolver
+
+ /**
+ * Obtains a [[ShuffleHandle]] to pass to tasks.
+ */
+ override def registerShuffle[K, V, C](
+ shuffleId: Int,
+ numMaps: Int,
+ dependency: ShuffleDependency[K, V, C]): ShuffleHandle = {
+ sortShuffleManager.registerShuffle(shuffleId, numMaps, dependency)
+ }
+
+ /**
+ * Get a reader for a range of reduce partitions (startPartition to endPartition-1, inclusive).
+ * Called on executors by reduce tasks.
+ */
+ override def getReader[K, C](
+ handle: ShuffleHandle,
+ startPartition: Int,
+ endPartition: Int,
+ context: TaskContext,
+ metrics: ShuffleReadMetricsReporter): ShuffleReader[K, C] = {
+
+ if (isArrowShuffle(handle)) {
+ new BlazeBlockStoreShuffleReader(
+ handle.asInstanceOf[BaseShuffleHandle[K, _, C]],
+ startPartition,
+ endPartition,
+ context,
+ metrics)
+ } else {
+ sortShuffleManager.getReader(handle, startPartition, endPartition, context, metrics)
+ }
+ }
+
+ override def getReader[K, C](
+ handle: ShuffleHandle,
+ startPartition: Int,
+ endPartition: Int,
+ context: TaskContext,
+ metrics: ShuffleReadMetricsReporter,
+ startMapId: Int,
+ endMapId: Int): ShuffleReader[K, C] = {
+
+ if (isArrowShuffle(handle)) {
+ new BlazeBlockStoreShuffleReader(
+ handle.asInstanceOf[BaseShuffleHandle[K, _, C]],
+ startPartition,
+ endPartition,
+ context,
+ metrics,
+ startMapId = Some(startMapId),
+ endMapId = Some(endMapId))
+ } else {
+ sortShuffleManager.getReader(
+ handle,
+ startPartition,
+ endPartition,
+ context,
+ metrics,
+ startMapId,
+ endMapId)
+ }
+ }
+
+ /** Get a writer for a given partition. Called on executors by map tasks. */
+ override def getWriter[K, V](
+ handle: ShuffleHandle,
+ mapId: Int,
+ context: TaskContext,
+ metrics: ShuffleWriteMetricsReporter): ShuffleWriter[K, V] = {
+ if (isArrowShuffle(handle)) {
+ new BlazeShuffleWriter(metrics)
+ } else {
+ sortShuffleManager.getWriter(handle, mapId, context, metrics)
+ }
+
+ }
+
+ /** Remove a shuffle's metadata from the ShuffleManager. */
+ override def unregisterShuffle(shuffleId: Int): Boolean = {
+ sortShuffleManager.unregisterShuffle(shuffleId)
+ }
+
+ /** Shut down this ShuffleManager. */
+ override def stop(): Unit = {
+ sortShuffleManager.stop()
+ }
+}
diff --git a/spark-extension-shims-spark3/src/main/scala/org/apache/spark/shuffle/uniffle/RssShuffleHandleWrapper.scala b/spark-extension-shims-spark241kwaiae/src/main/scala/org/apache/spark/sql/execution/blaze/shuffle/BlazeShuffleWriter.scala
similarity index 64%
copy from spark-extension-shims-spark3/src/main/scala/org/apache/spark/shuffle/uniffle/RssShuffleHandleWrapper.scala
copy to spark-extension-shims-spark241kwaiae/src/main/scala/org/apache/spark/sql/execution/blaze/shuffle/BlazeShuffleWriter.scala
index e8486b0a..f72eba7b 100644
--- a/spark-extension-shims-spark3/src/main/scala/org/apache/spark/shuffle/uniffle/RssShuffleHandleWrapper.scala
+++ b/spark-extension-shims-spark241kwaiae/src/main/scala/org/apache/spark/sql/execution/blaze/shuffle/BlazeShuffleWriter.scala
@@ -13,11 +13,9 @@
* See the License for the specific language governing permissions and
* limitations under the License.
*/
-package org.apache.spark.shuffle.uniffle
+package org.apache.spark.sql.execution.blaze.shuffle
-import org.apache.spark.shuffle.{BaseShuffleHandle, RssShuffleHandle}
+import org.apache.spark.shuffle.ShuffleWriteMetricsReporter
-class RssShuffleHandleWrapper[K, V, C](val rssShuffleHandleInfo: RssShuffleHandle[K, V, C])
- extends BaseShuffleHandle[K, V, C](
- rssShuffleHandleInfo.getShuffleId,
- rssShuffleHandleInfo.getDependency) {}
+class BlazeShuffleWriter[K, V](metrics: ShuffleWriteMetricsReporter)
+ extends BlazeShuffleWriterBase[K, V](metrics) {}
diff --git a/spark-extension-shims-spark241kwaiae/src/main/scala/org/apache/spark/sql/execution/blaze/shuffle/RssPartitionWriter.scala b/spark-extension-shims-spark241kwaiae/src/main/scala/org/apache/spark/sql/execution/blaze/shuffle/RssPartitionWriter.scala
new file mode 100644
index 00000000..a4b2e324
--- /dev/null
+++ b/spark-extension-shims-spark241kwaiae/src/main/scala/org/apache/spark/sql/execution/blaze/shuffle/RssPartitionWriter.scala
@@ -0,0 +1,86 @@
+/*
+ * Copyright 2022 The Blaze Authors
+ *
+ * Licensed 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.execution.blaze.shuffle
+
+import java.nio.ByteBuffer
+
+import org.apache.commons.lang3.reflect.FieldUtils
+import org.apache.spark.SparkEnv
+import org.apache.spark.TaskContext
+import org.apache.spark.internal.Logging
+import org.apache.spark.shuffle.ShuffleHandle
+import org.apache.spark.shuffle.ShuffleWriteMetricsReporter
+import org.apache.spark.shuffle.stream.StreamWriterWrapper
+
+import com.kuaishou.dataarch.shuffle.common.io.UnifiedDataInputBuffer
+import com.shade.org.apache.commons.lang3.reflect.MethodUtils
+
+class RssPartitionWriter(
+ handle: ShuffleHandle,
+ mapId: Int,
+ metrics: ShuffleWriteMetricsReporter,
+ numPartitions: Int)
+ extends RssPartitionWriterBase
+ with Logging {
+
+ private val tc = TaskContext.get()
+ private val streamWriter = SparkEnv.get.shuffleManager
+ .asInstanceOf[BlazeRssShuffleManager]
+ .getEffectiveManager(handle)
+ .getWriter(handle, mapId, tc, metrics)
+ private val realWriter = FieldUtils
+ .readField(streamWriter, "realWriter", true)
+ .asInstanceOf[StreamWriterWrapper[_, _]]
+ private val emptyKey = FieldUtils
+ .readField(streamWriter, "emptyKey", true)
+ .asInstanceOf[UnifiedDataInputBuffer]
+ private val partitionLengthMap = new Array[Long](numPartitions)
+
+ tc.addTaskCompletionListener(_ => this.close(true))
+ tc.addTaskFailureListener((_, _) => this.close(false))
+
+ override def write(partitionId: Int, buffer: ByteBuffer): Unit = {
+ val length = buffer.remaining()
+ val valueBuffer = FieldUtils
+ .readField(streamWriter, "value", true)
+ .asInstanceOf[UnifiedDataInputBuffer]
+
+ val bufferLen = buffer.limit() - buffer.position()
+ val bytes = new Array[Byte](bufferLen)
+ buffer.get(bytes)
+ valueBuffer.reset(bytes, length)
+
+ MethodUtils.invokeMethod(
+ realWriter,
+ true,
+ "append",
+ emptyKey,
+ valueBuffer,
+ Int.box(partitionId),
+ Int.box(0))
+
+ partitionLengthMap(partitionId) = length
+ metrics.incBytesWritten(length)
+ }
+
+ override def flush(): Unit = {
+ realWriter.spill()
+ }
+
+ override def close(success: Boolean): Unit = realWriter.close()
+
+ override def getPartitionLengthMap: Array[Long] = partitionLengthMap
+}
diff --git a/spark-extension-shims-spark241kwaiae/src/main/scala/org/apache/spark/sql/execution/joins/blaze/plan/NativeBroadcastJoinExec.scala b/spark-extension-shims-spark241kwaiae/src/main/scala/org/apache/spark/sql/execution/joins/blaze/plan/NativeBroadcastJoinExec.scala
new file mode 100644
index 00000000..224ab70a
--- /dev/null
+++ b/spark-extension-shims-spark241kwaiae/src/main/scala/org/apache/spark/sql/execution/joins/blaze/plan/NativeBroadcastJoinExec.scala
@@ -0,0 +1,59 @@
+/*
+ * Copyright 2022 The Blaze Authors
+ *
+ * Licensed 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.execution.joins.blaze.plan
+
+import org.apache.spark.sql.catalyst.expressions.Expression
+import org.apache.spark.sql.catalyst.plans.JoinType
+import org.apache.spark.sql.catalyst.plans.physical.Partitioning
+import org.apache.spark.sql.execution.SparkPlan
+import org.apache.spark.sql.execution.blaze.plan.BroadcastLeft
+import org.apache.spark.sql.execution.blaze.plan.BroadcastRight
+import org.apache.spark.sql.execution.blaze.plan.BroadcastSide
+import org.apache.spark.sql.execution.blaze.plan.NativeBroadcastJoinBase
+import org.apache.spark.sql.execution.joins
+import org.apache.spark.sql.execution.joins.HashJoin
+
+case class NativeBroadcastJoinExec(
+ override val left: SparkPlan,
+ override val right: SparkPlan,
+ override val outputPartitioning: Partitioning,
+ override val leftKeys: Seq[Expression],
+ override val rightKeys: Seq[Expression],
+ override val joinType: JoinType,
+ broadcastSide: BroadcastSide)
+ extends NativeBroadcastJoinBase(
+ left,
+ right,
+ outputPartitioning,
+ leftKeys,
+ rightKeys,
+ joinType,
+ broadcastSide)
+ with HashJoin {
+
+ override val condition: Option[Expression] = None
+
+ override val buildSide: joins.BuildSide = broadcastSide match {
+ case BroadcastLeft => joins.BuildLeft
+ case BroadcastRight => joins.BuildRight
+ }
+
+ override protected def rewriteKeyExprToLong(exprs: Seq[Expression]): Seq[Expression] =
+ HashJoin.rewriteKeyExpr(exprs)
+
+ override def withNewChildren(newChildren: Seq[SparkPlan]): SparkPlan =
+ copy(left = newChildren(0), right = newChildren(1))
+}
diff --git a/spark-extension-shims-spark241kwaiae/src/main/scala/org/apache/spark/sql/execution/joins/blaze/plan/NativeShuffledHashJoinExec.scala b/spark-extension-shims-spark241kwaiae/src/main/scala/org/apache/spark/sql/execution/joins/blaze/plan/NativeShuffledHashJoinExec.scala
new file mode 100644
index 00000000..6103ae23
--- /dev/null
+++ b/spark-extension-shims-spark241kwaiae/src/main/scala/org/apache/spark/sql/execution/joins/blaze/plan/NativeShuffledHashJoinExec.scala
@@ -0,0 +1,54 @@
+/*
+ * Copyright 2022 The Blaze Authors
+ *
+ * Licensed 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.execution.joins.blaze.plan
+
+import org.apache.spark.sql.catalyst.expressions.Expression
+import org.apache.spark.sql.catalyst.plans.JoinType
+import org.apache.spark.sql.execution.SparkPlan
+import org.apache.spark.sql.execution.blaze.plan.BuildLeft
+import org.apache.spark.sql.execution.blaze.plan.BuildRight
+import org.apache.spark.sql.execution.blaze.plan.BuildSide
+import org.apache.spark.sql.execution.blaze.plan.NativeShuffledHashJoinBase
+import org.apache.spark.sql.execution.joins
+import org.apache.spark.sql.execution.joins.HashJoin
+import org.apache.spark.sql.execution.joins.ShuffledHashJoinExec
+
+case class NativeShuffledHashJoinExec(
+ override val left: SparkPlan,
+ override val right: SparkPlan,
+ leftKeys: Seq[Expression],
+ rightKeys: Seq[Expression],
+ joinType: JoinType,
+ buildSide: BuildSide)
+ extends {
+ override val (output, outputPartitioning, outputOrdering) = {
+ val joinsBuildSide = buildSide match {
+ case BuildLeft => joins.BuildLeft
+ case BuildRight => joins.BuildRight
+ }
+ val bhj =
+ ShuffledHashJoinExec(leftKeys, rightKeys, joinType, joinsBuildSide, None, left, right)
+ (bhj.output, bhj.outputPartitioning, bhj.outputOrdering)
+ }
+ }
+ with NativeShuffledHashJoinBase(left, right, leftKeys, rightKeys, joinType, buildSide) {
+
+ override protected def rewriteKeyExprToLong(exprs: Seq[Expression]): Seq[Expression] =
+ HashJoin.rewriteKeyExpr(exprs)
+
+ override def withNewChildren(newChildren: Seq[SparkPlan]): SparkPlan =
+ copy(left = newChildren(0), right = newChildren(1))
+}
diff --git a/spark-extension-shims-spark241kwaiae/src/main/scala/org/apache/spark/sql/execution/ui/BlazeSQLPlanCollector.scala b/spark-extension-shims-spark241kwaiae/src/main/scala/org/apache/spark/sql/execution/ui/BlazeSQLPlanCollector.scala
new file mode 100644
index 00000000..cbcee766
--- /dev/null
+++ b/spark-extension-shims-spark241kwaiae/src/main/scala/org/apache/spark/sql/execution/ui/BlazeSQLPlanCollector.scala
@@ -0,0 +1,301 @@
+/*
+ * Copyright 2022 The Blaze Authors
+ *
+ * Licensed 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.execution.ui
+
+import com.fasterxml.jackson.annotation.JsonProperty
+import com.fasterxml.jackson.databind.ObjectMapper
+import com.fasterxml.jackson.databind.node.ObjectNode
+import org.apache.spark.internal.Logging
+import org.apache.spark.scheduler._
+import org.apache.spark.sql.execution.blaze.plan.NativeParquetScanExec
+import org.apache.spark.util.Utils
+import org.apache.spark.{KwaiSparkBasicMetrics, SparkContext}
+
+import java.text.NumberFormat
+import java.util.Locale
+import java.util.regex.Pattern
+import scala.beans.BeanProperty
+import scala.collection.mutable.ArrayBuffer
+import scala.util.Try
+
+class BlazeSQLPlanCollector(sc: SparkContext, statusStore: SQLAppStatusStore)
+ extends SparkListener
+ with Logging {
+
+ private val kafkaMetricsBroker =
+ sc.conf.get(
+ "spark.kwai.blazeOperatorMetrics.broker",
+ defaultValue =
+ "dataarch-bjdy-rs7110.idczw.hb1.kwaidc.com:9092,dataarch-bjdy-rs7111.idczw.hb1.kwaidc.com:9092,dataarch-bjdy-rs7112.idczw.hb1.kwaidc.com:9092,dataarch-bjdy-rs7113.idczw.hb1.kwaidc.com:9092,dataarch-bjdy-rs7114.idczw.hb1.kwaidc.com:9092")
+ private val kafkaMetricsTopic =
+ sc.conf.get(
+ "spark.kwai.blazeOperatorMetrics.topic",
+ defaultValue = "dataarch_blaze_operator_metrics")
+
+ private val producer: Option[KafkaClient] = Try(
+ KafkaClient(kafkaMetricsBroker, null, kafkaMetricsTopic)).toOption
+
+ private val objectMapper: ObjectMapper = new ObjectMapper()
+
+ def buildExecutionPlansInfo(executionId: Long): Unit = {
+ try {
+ val executionMetrics = statusStore.executionMetrics(executionId)
+ statusStore.planGraph(executionId).nodes.foreach { plan =>
+ BlazeSQLPlanCollector.plansInfo.append(buildPlanInfo(executionId, executionMetrics, plan))
+ }
+ } catch {
+ case e: Exception => // does not fail when proceed a metric error, ignore
+ }
+ }
+
+ def sendPlanMetrics2Kafka(): Unit = {
+ try {
+ val startTime = System.currentTimeMillis()
+ producer.foreach(
+ _.send(
+ key = sc.conf.getAppId,
+ objectMapper.writeValueAsString(
+ new BlazeApplicationInfo(sc, BlazeSQLPlanCollector.plansInfo.toArray))))
+ logInfo(
+ s"send app opretor metrics to kafka succ after ${System.currentTimeMillis() - startTime}ms")
+ } catch {
+ case e: Exception =>
+ logInfo(s"Blaze sendOperatorMetrics error: ${e.getMessage}")
+ producer.foreach(_.flush())
+ producer.foreach(_.close())
+ }
+ }
+
+ private def buildPlanInfo(
+ executionId: Long,
+ executionMetrics: Map[Long, String],
+ plan: SparkPlanGraphNode): BlazeOperatorInfo = {
+ val metricObject = objectMapper.createObjectNode()
+ metricObject.put("planIsOverride", 1)
+ plan.metrics.foreach { metric =>
+ try {
+ executionMetrics.get(metric.accumulatorId).foreach { rawMetricValue =>
+ parsePlanMetric(metric, rawMetricValue).foreach { case (metricName, metricValue) =>
+ metricObject.put(metricName, metricValue)
+ }
+ }
+ } catch {
+ case e: Exception =>
+ // does not fail when sending a metric error, ignore
+ logInfo(s"get metric ${metric.name} error because ${e.getMessage}")
+ }
+ }
+ new BlazeOperatorInfo(executionId, plan, metricObject)
+ }
+
+ private def parsePlanMetric(metric: SQLPlanMetric, metricValue: String): Seq[(String, Long)] = {
+ val metricName = metric.name
+ metric.metricType match {
+ case "sum" =>
+ val valueParser: (String => Number) = NumberFormat.getNumberInstance(Locale.US).parse
+ if (metric.name.equals("stageId")) {
+ Seq(
+ (metricName, valueParser(metricValue).longValue()),
+ (
+ "stage_num_tasks",
+ BlazeSQLPlanCollector.stageInfo
+ .getOrElse(valueParser(metricValue).intValue(), (-1L, 0L))
+ ._2),
+ (
+ "is_stage_failed",
+ BlazeSQLPlanCollector.stageInfo
+ .getOrElse(valueParser(metricValue).intValue(), (-1L, 0L))
+ ._1))
+ } else {
+ Seq((metricName, valueParser(metricValue).longValue()))
+ }
+
+ case "size" | "timing" | "nsTiming" =>
+ val sumMinMedMaxPattern = "(.*?)\\((.*?),(.*?),(.*?)\\)".r
+ val valueParser: (String => Double) = metric.metricType match {
+ case "size" => sizeStringToBytes
+ case "timing" | "nsTiming" => timeStringToMsDuration
+ }
+ val deleteSpaceValue = metricValue.replaceAll("\\s", "")
+ deleteSpaceValue match {
+ case sumMinMedMaxPattern(sum, min, med, max) =>
+ val frontPattern = metricName.replaceAll("\\(.*?\\)|\\s|(?i)total", "")
+ Seq(
+ (frontPattern + "_sum", valueParser(sum.split("\\s").mkString).longValue()),
+ (frontPattern + "_min", valueParser(min.split("\\s").mkString).longValue()),
+ (frontPattern + "_med", valueParser(med.split("\\s").mkString).longValue()),
+ (frontPattern + "_max", valueParser(max.split("\\s").mkString).longValue()))
+ case _ => Seq() // cannot parse, ignore
+ }
+
+ case "average" =>
+ val minMedMaxPattern = "\n\\((.*?), (.*?), (.*?)\\)".r
+ val valueParser: (String => Number) = NumberFormat.getNumberInstance(Locale.US).parse
+ metricValue match {
+ case minMedMaxPattern(max, med, min) =>
+ val frontPattern = metricName.replaceAll("\\(.*?\\)|\\s|(?i)total", "")
+ Seq(
+ (frontPattern + "_min", valueParser(min).longValue()),
+ (frontPattern + "_med", valueParser(med).longValue()),
+ (frontPattern + "_max", valueParser(max).longValue()))
+ case _ => Seq() // cannot parse, ignore
+ }
+
+ case _ => Seq() // unrecognized metric type, ignore
+ }
+ }
+
+ override def close(): Unit = {
+ producer.foreach(_.flush())
+ producer.foreach(_.close())
+ }
+
+ def sizeStringToBytes(size: String): Double = {
+ val EB = 1L << 60
+ val PB = 1L << 50
+ val TB = 1L << 40
+ val GB = 1L << 30
+ val MB = 1L << 20
+ val KB = 1L << 10
+ val B = 1L
+
+ val upper = size.toUpperCase(Locale.ROOT).trim()
+ try {
+ val matcher = Pattern.compile("(\\d+(?:\\.\\d+)?)([A-Z]+)?").matcher(upper)
+ if (matcher.matches()) {
+ val value = matcher.group(1).toDouble
+ val suffix = matcher.group(2)
+ suffix match {
+ case "EB" => value * EB
+ case "PB" => value * PB
+ case "TB" => value * TB
+ case "GB" => value * GB
+ case "MB" => value * MB
+ case "KB" => value * KB
+ case "B" => value * B
+ case _ => value
+ }
+ } else {
+ throw new NumberFormatException("Failed to parse byte string: " + size)
+ }
+ } catch {
+ case _ => throw new NumberFormatException("Failed to parse byte string: " + size)
+ }
+ }
+
+ def timeStringToMsDuration(time: String): Double = {
+ val ms = 1L
+ val s = 1000 * ms
+ val m = 60 * s
+ val min = m
+ val h = 60 * m
+ val d = 24 * h
+
+ val lower = time.toLowerCase(Locale.ROOT).trim()
+ try {
+ val matcher = Pattern.compile("(\\d+(?:\\.\\d+)?)([a-z]+)?").matcher(lower)
+ if (matcher.matches()) {
+ val value = matcher.group(1).toDouble
+ val suffix = matcher.group(2)
+ suffix match {
+ case "ms" => value * ms
+ case "s" => value * s
+ case "m" => value * m
+ case "min" => value * min
+ case "h" => value * h
+ case "d" => value * d
+ case _ => value
+ }
+ } else {
+ throw new NumberFormatException("Failed to parse time string: " + time)
+ }
+ } catch {
+ case _ => {
+ throw new NumberFormatException("Failed to parse time string: " + time)
+ }
+ }
+ }
+
+}
+
+class BlazeApplicationInfo(sc: SparkContext, info: Array[BlazeOperatorInfo])
+ extends KwaiSparkBasicMetrics(sc) {
+ @BeanProperty
+ @JsonProperty("blazeJarVersion")
+ val blazeJarVersion: String = sc.conf.get("spark.blaze.jar", "nonblaze")
+ @BeanProperty
+ @JsonProperty("operatorsInfo")
+ val operatorsInfo: Array[BlazeOperatorInfo] = info
+}
+
+class BlazeOperatorInfo(eid: Long, plan: SparkPlanGraphNode, execMetric: ObjectNode) {
+ @BeanProperty
+ @JsonProperty("executionId")
+ val executionId: Long = eid
+ @BeanProperty
+ @JsonProperty("planId")
+ val planId: Long = plan.id
+ @BeanProperty
+ @JsonProperty("planName")
+ val planName: String = if (plan.name.contains("NativeParquetScan")) {
+ Utils.getSimpleName(NativeParquetScanExec.getClass)
+ } else {
+ plan.name
+ }
+ @BeanProperty
+ @JsonProperty("planMetric")
+ val planMetric: ObjectNode = execMetric
+}
+
+object BlazeSQLPlanCollector {
+ var stageInfo = scala.collection.mutable.Map[Int, (Long, Long)]()
+ var plansInfo = new ArrayBuffer[BlazeOperatorInfo]()
+}
+
+class BlazePlanListener(sparkContext: SparkContext, statusStore: SQLAppStatusStore)
+ extends SparkListener
+ with Logging {
+
+ val metricsCollector: BlazeSQLPlanCollector =
+ new BlazeSQLPlanCollector(sparkContext, statusStore)
+
+ override def onOtherEvent(event: SparkListenerEvent): Unit = event match {
+ case e: SparkListenerSQLExecutionEnd => onExecutionEnd(e)
+ case _ => // Ignore
+ }
+
+ override def onStageCompleted(stageCompleted: SparkListenerStageCompleted): Unit = {
+ val info = stageCompleted.stageInfo
+ BlazeSQLPlanCollector.stageInfo.put(
+ info.stageId,
+ (info.getStatusString.equals("failed").compare(false), info.numTasks))
+ }
+
+ private def onExecutionEnd(event: SparkListenerSQLExecutionEnd): Unit = {
+ metricsCollector.buildExecutionPlansInfo(event.executionId)
+ }
+
+ override def onApplicationEnd(applicationEnd: SparkListenerApplicationEnd): Unit = {
+ metricsCollector.sendPlanMetrics2Kafka()
+ }
+
+ override def close(): Unit = {
+ sparkContext.listenerBus.removeListener(this)
+ metricsCollector.close()
+ BlazeSQLPlanCollector.stageInfo.clear()
+ BlazeSQLPlanCollector.plansInfo.clear()
+ }
+}
diff --git a/spark-extension-shims-spark3/src/main/scala/org/apache/spark/shuffle/uniffle/RssShuffleHandleWrapper.scala b/spark-extension-shims-spark3/src/main/scala/org/apache/spark/shuffle/uniffle/RssShuffleHandleWrapper.scala
index e8486b0a..2028d5a1 100644
--- a/spark-extension-shims-spark3/src/main/scala/org/apache/spark/shuffle/uniffle/RssShuffleHandleWrapper.scala
+++ b/spark-extension-shims-spark3/src/main/scala/org/apache/spark/shuffle/uniffle/RssShuffleHandleWrapper.scala
@@ -20,4 +20,8 @@ import org.apache.spark.shuffle.{BaseShuffleHandle, RssShuffleHandle}
class RssShuffleHandleWrapper[K, V, C](val rssShuffleHandleInfo: RssShuffleHandle[K, V, C])
extends BaseShuffleHandle[K, V, C](
rssShuffleHandleInfo.getShuffleId,
- rssShuffleHandleInfo.getDependency) {}
+ numMaps = 0,
+ rssShuffleHandleInfo.getDependency) {
+
+ throw new UnsupportedOperationException("uniffle is not supported yet")
+}
diff --git a/spark-extension-shims-spark3/src/main/scala/org/apache/spark/sql/blaze/ShimsImpl.scala b/spark-extension-shims-spark3/src/main/scala/org/apache/spark/sql/blaze/ShimsImpl.scala
index 7bb9f2f6..44ce358e 100644
--- a/spark-extension-shims-spark3/src/main/scala/org/apache/spark/sql/blaze/ShimsImpl.scala
+++ b/spark-extension-shims-spark3/src/main/scala/org/apache/spark/sql/blaze/ShimsImpl.scala
@@ -30,6 +30,8 @@ import org.apache.spark.scheduler.MapStatus
import org.apache.spark.shuffle.IndexShuffleBlockResolver
import org.apache.spark.shuffle.ShuffleHandle
import org.apache.spark.shuffle.ShuffleWriteMetricsReporter
+import org.apache.spark.shuffle.sort.MapInfo
+import org.apache.spark.shuffle.BaseShuffleHandle
import org.apache.spark.sql.SQLContext
import org.apache.spark.sql.SparkSession
import org.apache.spark.sql.blaze.BlazeConverters.ForceNativeExecutionWrapperBase
@@ -109,9 +111,12 @@ import org.apache.spark.sql.hive.execution.InsertIntoHiveTable
import org.apache.spark.sql.types.DataType
import org.apache.spark.sql.types.IntegerType
import org.apache.spark.sql.types.StringType
+import org.apache.spark.sql.SparkSessionExtensions
import org.apache.spark.storage.BlockManagerId
import org.apache.spark.storage.FileSegment
import org.apache.spark.sql.execution.blaze.shuffle.uniffle.BlazeUniffleShuffleManager
+import org.apache.spark.storage.ShuffleDataBlockId
+import org.apache.spark.util.ExternalBlockStoreUtils
import org.blaze.{protobuf => pb, sparkver}
class ShimsImpl extends Shims with Logging {
@@ -144,10 +149,12 @@ class ShimsImpl extends Shims with Logging {
}
}
- @sparkver("3.0 / 3.1")
- override def initExtension(): Unit = {
- if (BlazeConf.FORCE_SHUFFLED_HASH_JOIN.booleanConf()) {
- logWarning(s"${BlazeConf.FORCE_SHUFFLED_HASH_JOIN.key} is not supported in $shimVersion")
+ override def onApplyingExtension(extension: SparkSessionExtensions): Unit = {
+ @sparkver("3.0 / 3.1")
+ val _impl: Unit = {
+ if (BlazeConf.FORCE_SHUFFLED_HASH_JOIN.booleanConf()) {
+ logWarning(s"${BlazeConf.FORCE_SHUFFLED_HASH_JOIN.key} is not supported in $shimVersion")
+ }
}
}
@@ -374,9 +381,11 @@ class ShimsImpl extends Shims with Logging {
exec
}
- override def getRDDShuffleReadFull(rdd: RDD[_]): Boolean = true
+ override def getRDDShuffleReadFull(rdd: RDD[_]): Boolean = rdd.shuffleReadFull
- override def setRDDShuffleReadFull(rdd: RDD[_], shuffleReadFull: Boolean): Unit = {}
+ override def setRDDShuffleReadFull(rdd: RDD[_], shuffleReadFull: Boolean): Unit = {
+ rdd.shuffleReadFull = shuffleReadFull
+ }
override def createFileSegment(
file: File,
@@ -401,7 +410,44 @@ class ShimsImpl extends Shims with Logging {
partitionLengths,
checksums,
tempDataFile)
- MapStatus.apply(SparkEnv.get.blockManager.shuffleServerId, partitionLengths, mapId)
+
+ // shuffle write on hdfs
+ val blockManager = SparkEnv.get.blockManager
+ val handle = dep.shuffleHandle.asInstanceOf[BaseShuffleHandle[_, _, _]]
+ val output = shuffleBlockResolver.getDataFile(dep.shuffleId, mapId.toInt)
+ val mapInfo = new MapInfo(partitionLengths, partitionLengths.map(_ => 0L))
+ val totalPartitionLength = mapInfo.lengths.sum
+ var hasExternalData = false
+
+ if (ExternalBlockStoreUtils.writeRemoteEnabled(
+ SparkEnv.get.conf,
+ totalPartitionLength,
+ handle.numMaps)) {
+ if (dataSize != 0) {
+ logInfo(s"KwaiShuffle: Start to write remote, shuffle block id ${ShuffleDataBlockId(
+ dep.shuffleId,
+ mapId.toInt,
+ 0).name}, file path is ${output.getPath}/${output.getName}")
+ context.taskMetrics.externalMetrics.writeRemoteShuffle.setValue(1L)
+ val dataBlockId = new ShuffleDataBlockId(
+ dep.shuffleId,
+ mapId.toInt,
+ IndexShuffleBlockResolver.NOOP_REDUCE_ID)
+ val indexFile = shuffleBlockResolver.getIndexFile(dep.shuffleId, mapId.toInt)
+ blockManager.externalBlockStore.externalBlockManager.get.writeExternalShuffleFile(
+ dep.shuffleId,
+ TaskContext.get.partitionId, // in spark3.5, use partition id other than map id
+ dataBlockId,
+ indexFile,
+ output)
+ hasExternalData = true
+ }
+ }
+
+ val mapStatus =
+ getMapStatus(SparkEnv.get.blockManager.shuffleServerId, partitionLengths, mapId)
+ mapStatus.hasExternal = hasExternalData
+ mapStatus
}
@sparkver("3.0 / 3.1")
@@ -419,7 +465,7 @@ class ShimsImpl extends Shims with Logging {
mapId,
partitionLengths,
tempDataFile)
- MapStatus.apply(SparkEnv.get.blockManager.shuffleServerId, partitionLengths, mapId)
+ getMapStatus(SparkEnv.get.blockManager.shuffleServerId, partitionLengths, mapId)
}
override def getRssPartitionWriter(
@@ -432,7 +478,7 @@ class ShimsImpl extends Shims with Logging {
shuffleServerId: BlockManagerId,
partitionLengthMap: Array[Long],
mapId: Long): MapStatus =
- MapStatus.apply(shuffleServerId, partitionLengthMap, mapId)
+ MapStatus.apply(shuffleServerId, partitionLengthMap, partitionLengthMap.map(_ => 0L), mapId)
override def getShuffleWriteExec(
input: pb.PhysicalPlanNode,
@@ -569,6 +615,7 @@ class ShimsImpl extends Shims with Logging {
val shuffleHandle = shuffledRDD.dependency.shuffleHandle
val inputRDD = executeNative(child)
+ val inputMetrics = inputRDD.metrics
val nativeShuffle = getUnderlyingNativePlan(child).asInstanceOf[NativeShuffleExchangeExec]
val nativeSchema: pb.Schema = nativeShuffle.nativeSchema
@@ -592,7 +639,7 @@ class ShimsImpl extends Shims with Logging {
shuffledRDD.partitions,
shuffledRDD.partitioner,
new OneToOneDependency(shuffledRDD) :: Nil,
- true,
+ shuffledRDD.shuffleReadFull,
(partition, taskContext) => {
// use reflection to get partitionSpec because ShuffledRowRDDPartition is private
@@ -976,7 +1023,6 @@ class ShimsImpl extends Shims with Logging {
e: Expression,
isPruningExpr: Boolean,
fallback: Expression => pb.PhysicalExprNode): Option[pb.PhysicalExprNode] = None
-
}
case class ForceNativeExecutionWrapper(override val child: SparkPlan)
diff --git a/spark-extension-shims-spark3/src/main/scala/org/apache/spark/sql/execution/blaze/kwai/KwaiPrivilegedHDFSBlockManager.scala b/spark-extension-shims-spark3/src/main/scala/org/apache/spark/sql/execution/blaze/kwai/KwaiPrivilegedHDFSBlockManager.scala
new file mode 100644
index 00000000..7f96eed5
--- /dev/null
+++ b/spark-extension-shims-spark3/src/main/scala/org/apache/spark/sql/execution/blaze/kwai/KwaiPrivilegedHDFSBlockManager.scala
@@ -0,0 +1,51 @@
+/*
+ * Copyright 2022 The Blaze Authors
+ *
+ * Licensed 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.execution.blaze.kwai
+
+import java.security.PrivilegedExceptionAction
+
+import org.apache.spark.sql.blaze.NativeHelper
+import org.apache.spark.storage.BlockManager
+import org.apache.spark.storage.HDFSBlockManager
+import org.apache.spark.SparkConf
+import org.apache.spark.internal.Logging
+import org.apache.spark.storage.ExternalBlockStore
+
+class KwaiPrivilegedHDFSBlockManager extends HDFSBlockManager {
+ override def init(blockManager: BlockManager): Unit = {
+ val initHDFSBlockManager = () => {
+ super.init(blockManager)
+ }
+
+ NativeHelper.currentUser.doAs(new PrivilegedExceptionAction[Unit] {
+ override def run(): Unit = initHDFSBlockManager()
+ })
+ }
+}
+
+object KwaiPrivilegedHDFSBlockManager extends Logging {
+ def setup(conf: SparkConf): Unit = {
+ val externalBlockManagerClassName = conf
+ .getOption(ExternalBlockStore.BLOCK_MANAGER_NAME)
+ .getOrElse(ExternalBlockStore.DEFAULT_BLOCK_MANAGER_NAME)
+ val thisClassName = classOf[KwaiPrivilegedHDFSBlockManager].getName
+
+ if (externalBlockManagerClassName == classOf[HDFSBlockManager].getName) {
+ logInfo(s"replacing external block manager to $thisClassName")
+ conf.set(ExternalBlockStore.BLOCK_MANAGER_NAME, thisClassName)
+ }
+ }
+}
diff --git a/spark-extension-shims-spark3/src/main/scala/org/apache/spark/sql/execution/blaze/shuffle/BlazeRssShuffleManagerBase.scala b/spark-extension-shims-spark3/src/main/scala/org/apache/spark/sql/execution/blaze/shuffle/BlazeRssShuffleManagerBase.scala
index f4574aa3..9e2ce10a 100644
--- a/spark-extension-shims-spark3/src/main/scala/org/apache/spark/sql/execution/blaze/shuffle/BlazeRssShuffleManagerBase.scala
+++ b/spark-extension-shims-spark3/src/main/scala/org/apache/spark/sql/execution/blaze/shuffle/BlazeRssShuffleManagerBase.scala
@@ -24,6 +24,7 @@ import org.blaze.sparkver
abstract class BlazeRssShuffleManagerBase(_conf: SparkConf) extends ShuffleManager with Logging {
override def registerShuffle[K, V, C](
shuffleId: Int,
+ numMaps: Int,
dependency: ShuffleDependency[K, V, C]): ShuffleHandle
override def unregisterShuffle(shuffleId: Int): Boolean
diff --git a/spark-extension-shims-spark3/src/main/scala/org/apache/spark/sql/execution/blaze/shuffle/BlazeShuffleManager.scala b/spark-extension-shims-spark3/src/main/scala/org/apache/spark/sql/execution/blaze/shuffle/BlazeShuffleManager.scala
index 169e31f8..50c42f7a 100644
--- a/spark-extension-shims-spark3/src/main/scala/org/apache/spark/sql/execution/blaze/shuffle/BlazeShuffleManager.scala
+++ b/spark-extension-shims-spark3/src/main/scala/org/apache/spark/sql/execution/blaze/shuffle/BlazeShuffleManager.scala
@@ -46,8 +46,17 @@ class BlazeShuffleManager(conf: SparkConf) extends ShuffleManager with Logging {
*/
override def registerShuffle[K, V, C](
shuffleId: Int,
+ numMaps: Int,
dependency: ShuffleDependency[K, V, C]): ShuffleHandle = {
- sortShuffleManager.registerShuffle(shuffleId, dependency)
+
+ val registerShuffleMethod = sortShuffleManager.getClass.getMethod(
+ "registerShuffle",
+ classOf[Int],
+ classOf[Int],
+ classOf[ShuffleDependency[K, V, C]])
+ registerShuffleMethod
+ .invoke(sortShuffleManager, Int.box(shuffleId), Int.box(numMaps), dependency)
+ .asInstanceOf[ShuffleHandle]
}
@sparkver("3.2 / 3.3 / 3.4 / 3.5")
diff --git a/spark-extension-shims-spark3/src/main/scala/org/apache/spark/sql/execution/blaze/shuffle/celeborn/BlazeCelebornShuffleManager.scala b/spark-extension-shims-spark3/src/main/scala/org/apache/spark/sql/execution/blaze/shuffle/celeborn/BlazeCelebornShuffleManager.scala
index e5002966..8ceb6536 100644
--- a/spark-extension-shims-spark3/src/main/scala/org/apache/spark/sql/execution/blaze/shuffle/celeborn/BlazeCelebornShuffleManager.scala
+++ b/spark-extension-shims-spark3/src/main/scala/org/apache/spark/sql/execution/blaze/shuffle/celeborn/BlazeCelebornShuffleManager.scala
@@ -30,8 +30,9 @@ class BlazeCelebornShuffleManager(conf: SparkConf, isDriver: Boolean)
override def registerShuffle[K, V, C](
shuffleId: Int,
+ numMaps: Int,
dependency: ShuffleDependency[K, V, C]): ShuffleHandle = {
- celebornShuffleManager.registerShuffle(shuffleId, dependency)
+ celebornShuffleManager.registerShuffle(shuffleId, numMaps, dependency)
}
override def unregisterShuffle(shuffleId: Int): Boolean = {
diff --git a/spark-extension-shims-spark3/src/main/scala/org/apache/spark/sql/execution/blaze/shuffle/uniffle/BlazeUniffleShuffleManager.scala b/spark-extension-shims-spark3/src/main/scala/org/apache/spark/sql/execution/blaze/shuffle/uniffle/BlazeUniffleShuffleManager.scala
index 96754cb8..6bbeb694 100644
--- a/spark-extension-shims-spark3/src/main/scala/org/apache/spark/sql/execution/blaze/shuffle/uniffle/BlazeUniffleShuffleManager.scala
+++ b/spark-extension-shims-spark3/src/main/scala/org/apache/spark/sql/execution/blaze/shuffle/uniffle/BlazeUniffleShuffleManager.scala
@@ -27,8 +27,9 @@ class BlazeUniffleShuffleManager(conf: SparkConf, isDriver: Boolean)
private val uniffleShuffleManager: RssShuffleManager = new RssShuffleManager(conf, isDriver);
override def registerShuffle[K, V, C](
shuffleId: Int,
+ numMaps: Int,
dependency: ShuffleDependency[K, V, C]): ShuffleHandle = {
- val handle = uniffleShuffleManager.registerShuffle(shuffleId, dependency)
+ val handle = uniffleShuffleManager.registerShuffle(shuffleId, numMaps, dependency)
new RssShuffleHandleWrapper(handle.asInstanceOf[RssShuffleHandle[K, V, C]])
}
diff --git a/spark-extension-shims-spark3/src/test/scala/org/apache/spark/sql/blaze/BaseBlazeSQLSuite.scala b/spark-extension-shims-spark3/src/test/scala/org/apache/spark/sql/blaze/BaseBlazeSQLSuite.scala
index 0ad73598..480ca856 100644
--- a/spark-extension-shims-spark3/src/test/scala/org/apache/spark/sql/blaze/BaseBlazeSQLSuite.scala
+++ b/spark-extension-shims-spark3/src/test/scala/org/apache/spark/sql/blaze/BaseBlazeSQLSuite.scala
@@ -15,8 +15,10 @@
*/
package org.apache.spark.sql.blaze
-import org.apache.spark.SparkConf
import org.apache.spark.sql.test.SharedSparkSession
+import org.apache.spark.SparkConf
+
+// test is not available in kwai branch
trait BaseBlazeSQLSuite extends SharedSparkSession {
@@ -29,5 +31,4 @@ trait BaseBlazeSQLSuite extends SharedSparkSession {
.set("spark.memory.offHeap.enabled", "false")
.set("spark.blaze.enable", "true")
}
-
}
diff --git a/spark-extension-shims-spark3/src/test/scala/org/apache/spark/sql/blaze/BlazeFunctionSuite.scala b/spark-extension-shims-spark3/src/test/scala/org/apache/spark/sql/blaze/BlazeFunctionSuite.scala
index d41dc03e..8f313f13 100644
--- a/spark-extension-shims-spark3/src/test/scala/org/apache/spark/sql/blaze/BlazeFunctionSuite.scala
+++ b/spark-extension-shims-spark3/src/test/scala/org/apache/spark/sql/blaze/BlazeFunctionSuite.scala
@@ -15,8 +15,6 @@
*/
package org.apache.spark.sql.blaze
-import org.apache.spark.sql.Row
-
class BlazeFunctionSuite extends org.apache.spark.sql.QueryTest with BaseBlazeSQLSuite {
test("sum function with float input") {
@@ -29,31 +27,4 @@ class BlazeFunctionSuite extends org.apache.spark.sql.QueryTest with BaseBlazeSQ
}
}
}
-
- test("sha2 function") {
- withTable("t1") {
- sql("create table t1 using parquet as select 'spark' as c1, '3.x' as version")
- val functions =
- """
- |select
- | sha2(concat(c1, version), 256) as sha0,
- | sha2(concat(c1, version), 256) as sha256,
- | sha2(concat(c1, version), 224) as sha224,
- | sha2(concat(c1, version), 384) as sha384,
- | sha2(concat(c1, version), 512) as sha512
- |from t1
- |""".stripMargin
- val df = sql(functions)
- checkAnswer(
- df,
- Seq(
- Row(
- "562d20689257f3f3a04ee9afb86d0ece2af106cf6c6e5e7d266043088ce5fbc0",
- "562d20689257f3f3a04ee9afb86d0ece2af106cf6c6e5e7d266043088ce5fbc0",
- "d0c8e9ccd5c7b3fdbacd2cfd6b4d65ca8489983b5e8c7c64cd77b634",
- "77c1199808053619c29e9af2656e1ad2614772f6ea605d5757894d6aec2dfaf34ff6fd662def3b79e429e9ae5ecbfed1",
- "c4e27d35517ca62243c1f322d7922dac175830be4668e8a1cf3befdcd287bb5b6f8c5f041c9d89e4609c8cfa242008c7c7133af1685f57bac9052c1212f1d089")))
- }
- }
-
}
diff --git a/spark-extension/src/main/scala/org/apache/spark/sql/blaze/BlazeConverters.scala b/spark-extension/src/main/scala/org/apache/spark/sql/blaze/BlazeConverters.scala
index 29a39c6c..c9af4188 100644
--- a/spark-extension/src/main/scala/org/apache/spark/sql/blaze/BlazeConverters.scala
+++ b/spark-extension/src/main/scala/org/apache/spark/sql/blaze/BlazeConverters.scala
@@ -135,7 +135,7 @@ object BlazeConverters extends Logging {
val enableScanParquet: Boolean =
SparkEnv.get.conf.getBoolean("spark.blaze.enable.scan.parquet", defaultValue = true)
val enableScanOrc: Boolean =
- SparkEnv.get.conf.getBoolean("spark.blaze.enable.scan.orc", defaultValue = true)
+ SparkEnv.get.conf.getBoolean("spark.blaze.enable.scan.orc", defaultValue = false)
import org.apache.spark.sql.catalyst.plans._
import org.apache.spark.sql.catalyst.optimizer._
@@ -313,16 +313,22 @@ object BlazeConverters extends Logging {
getShuffleOrigin(exec))
}
+ @sparkver("3.0 / 3.1 / 3.2 / 3.3 / 3.4 / 3.5")
+ def getIsSkewJoinFromSMJ(exec: SortMergeJoinExec): Boolean = exec.isSkewJoin
+
+ @sparkver("241kwaiae")
+ def getIsSkewJoinFromSMJ(exec: SortMergeJoinExec): Boolean = false
+
@sparkver(" 3.2 / 3.3 / 3.4 / 3.5")
def getIsSkewJoinFromSHJ(exec: ShuffledHashJoinExec): Boolean = exec.isSkewJoin
- @sparkver("3.0 / 3.1")
+ @sparkver("3.0 / 3.1 / 241kwaiae")
def getIsSkewJoinFromSHJ(exec: ShuffledHashJoinExec): Boolean = false
@sparkver("3.1 / 3.2 / 3.3 / 3.4 / 3.5")
def getShuffleOrigin(exec: ShuffleExchangeExec): Option[Any] = Some(exec.shuffleOrigin)
- @sparkver("3.0")
+ @sparkver("3.0 / 241kwaiae")
def getShuffleOrigin(exec: ShuffleExchangeExec): Option[Any] = None
def convertFileSourceScanExec(exec: FileSourceScanExec): SparkPlan = {
@@ -353,6 +359,9 @@ object BlazeConverters extends Logging {
relation.fileFormat match {
case p if p.getClass.getName.endsWith("ParquetFileFormat") =>
assert(enableScanParquet)
+ assert(
+ !p.getClass.getName.endsWith("HoodieFileGroupReaderBasedParquetFileFormat"),
+ "Hudi MOR table not supported")
addRenameColumnsExec(Shims.get.createNativeParquetScanExec(exec))
case p if p.getClass.getName.endsWith("OrcFileFormat") =>
assert(enableScanOrc)
@@ -419,7 +428,7 @@ object BlazeConverters extends Logging {
exec.condition,
exec.left,
exec.right,
- exec.isSkewJoin)
+ getIsSkewJoinFromSMJ(exec))
logDebug(
s"Converting SortMergeJoinExec (with forceShuffledHashJoin): ${Shims.get.simpleStringWithNodeId(exec)}")
logDebug(s" leftKeys: $leftKeys")
@@ -455,7 +464,7 @@ object BlazeConverters extends Logging {
exec.condition,
exec.left,
exec.right,
- exec.isSkewJoin)
+ getIsSkewJoinFromSMJ(exec))
logDebug(s"Converting SortMergeJoinExec: ${Shims.get.simpleStringWithNodeId(exec)}")
logDebug(s" leftKeys: $leftKeys")
logDebug(s" rightKeys: $rightKeys")
diff --git a/spark-extension/src/main/scala/org/apache/spark/sql/blaze/BlazeSparkSessionExtension.scala b/spark-extension/src/main/scala/org/apache/spark/sql/blaze/BlazeSparkSessionExtension.scala
index 45d75b94..6db25c84 100644
--- a/spark-extension/src/main/scala/org/apache/spark/sql/blaze/BlazeSparkSessionExtension.scala
+++ b/spark-extension/src/main/scala/org/apache/spark/sql/blaze/BlazeSparkSessionExtension.scala
@@ -35,7 +35,7 @@ class BlazeSparkSessionExtension extends (SparkSessionExtensions => Unit) with L
logInfo("org.apache.spark.BlazeSparkSessionExtension enabled")
assert(BlazeSparkSessionExtension.blazeEnabledKey != null)
- Shims.get.onApplyingExtension()
+ Shims.get.onApplyingExtension(extensions)
extensions.injectColumnar(sparkSession => {
BlazeColumnarOverrides(sparkSession)
diff --git a/spark-extension/src/main/scala/org/apache/spark/sql/blaze/NativeConverters.scala b/spark-extension/src/main/scala/org/apache/spark/sql/blaze/NativeConverters.scala
index 35c66055..4a57d474 100644
--- a/spark-extension/src/main/scala/org/apache/spark/sql/blaze/NativeConverters.scala
+++ b/spark-extension/src/main/scala/org/apache/spark/sql/blaze/NativeConverters.scala
@@ -19,18 +19,15 @@ import java.io.ByteArrayInputStream
import java.io.ByteArrayOutputStream
import java.io.ObjectInputStream
import java.io.ObjectOutputStream
-
import scala.collection.JavaConverters._
import scala.collection.mutable
import scala.math.max
import scala.math.min
-
import org.apache.arrow.c.CDataDictionaryProvider
import org.apache.arrow.vector.VectorSchemaRoot
import org.apache.arrow.vector.ipc.ArrowStreamWriter
import org.apache.commons.lang3.reflect.FieldUtils
import org.apache.commons.lang3.reflect.MethodUtils
-
import com.google.protobuf.ByteString
import org.apache.spark.SparkEnv
import org.blaze.{protobuf => pb}
@@ -48,7 +45,8 @@ import org.apache.spark.sql.catalyst.plans.LeftOuter
import org.apache.spark.sql.catalyst.plans.LeftSemi
import org.apache.spark.sql.catalyst.plans.RightOuter
import org.apache.spark.sql.catalyst.InternalRow
-import org.apache.spark.sql.catalyst.optimizer.NormalizeNaNAndZero
+import org.blaze.protobuf.ScalarFunction
+// import org.apache.spark.sql.catalyst.optimizer.NormalizeNaNAndZero
import org.apache.spark.sql.catalyst.plans.ExistenceJoin
import org.apache.spark.sql.execution.blaze.plan.Util
import org.apache.spark.sql.execution.ExecSubqueryExpression
@@ -339,13 +337,16 @@ object NativeConverters extends Logging {
convertExprWithFallback(
sparkExpr,
isPruningExpr = true,
- { _ =>
- buildExprNode(
- _.setColumn(
- pb.PhysicalColumn
+ { e => // convert unsupported expression to a placeholder
+ pb.PhysicalExprNode
+ .newBuilder()
+ .setScalarFunction(
+ pb.PhysicalScalarFunctionNode
.newBuilder()
- .setName("!__unsupported_pruning_expr__")
- .setIndex(Int.MaxValue)))
+ .setFun(pb.ScalarFunction.SparkExtFunctions)
+ .setName("Placeholder")
+ .setReturnType(NativeConverters.convertDataType(e.dataType)))
+ .build()
})
}
@@ -957,9 +958,9 @@ object NativeConverters extends Logging {
.apply(precision, IntegerType) :: Literal.apply(scale, IntegerType) :: Nil
buildExtScalarFunction("CheckOverflow", args, DecimalType(precision, scale))
- case e: NormalizeNaNAndZero
- if e.dataType.isInstanceOf[FloatType] || e.dataType.isInstanceOf[DoubleType] =>
- buildExtScalarFunction("NormalizeNanAndZero", e.children, e.dataType)
+ // case e: NormalizeNaNAndZero
+ // if e.dataType.isInstanceOf[FloatType] || e.dataType.isInstanceOf[DoubleType] =>
+ // buildExtScalarFunction("NormalizeNanAndZero", e.children, e.dataType)
case e: CreateArray => buildExtScalarFunction("MakeArray", e.children, e.dataType)
diff --git a/spark-extension/src/main/scala/org/apache/spark/sql/blaze/Shims.scala b/spark-extension/src/main/scala/org/apache/spark/sql/blaze/Shims.scala
index a6677e57..bb8a3994 100644
--- a/spark-extension/src/main/scala/org/apache/spark/sql/blaze/Shims.scala
+++ b/spark-extension/src/main/scala/org/apache/spark/sql/blaze/Shims.scala
@@ -16,6 +16,7 @@
package org.apache.spark.sql.blaze
import java.io.File
+
import org.apache.spark.ShuffleDependency
import org.apache.spark.TaskContext
import org.apache.spark.SparkContext
@@ -49,6 +50,7 @@ import org.apache.spark.sql.SparkSession
import org.apache.spark.sql.catalyst.InternalRow
import org.apache.spark.sql.catalyst.catalog.CatalogTable
import org.apache.spark.sql.execution.datasources.PartitionedFile
+import org.apache.spark.sql.SparkSessionExtensions
import org.apache.spark.storage.BlockManagerId
import org.apache.spark.storage.FileSegment
@@ -58,7 +60,7 @@ abstract class Shims {
def initExtension(): Unit = {}
- def onApplyingExtension(): Unit = {}
+ def onApplyingExtension(extension: SparkSessionExtensions): Unit
def createConvertToNativeExec(child: SparkPlan): ConvertToNativeBase
diff --git a/spark-extension/src/main/scala/org/apache/spark/sql/execution/blaze/plan/NativeBroadcastExchangeBase.scala b/spark-extension/src/main/scala/org/apache/spark/sql/execution/blaze/plan/NativeBroadcastExchangeBase.scala
index eb64236f..3fccb55a 100644
--- a/spark-extension/src/main/scala/org/apache/spark/sql/execution/blaze/plan/NativeBroadcastExchangeBase.scala
+++ b/spark-extension/src/main/scala/org/apache/spark/sql/execution/blaze/plan/NativeBroadcastExchangeBase.scala
@@ -64,6 +64,7 @@ import org.apache.spark.sql.execution.metric.SQLMetric
import org.apache.spark.sql.execution.metric.SQLMetrics
import org.apache.spark.sql.types.BinaryType
import org.blaze.{protobuf => pb}
+import org.blaze.sparkver
abstract class NativeBroadcastExchangeBase(mode: BroadcastMode, override val child: SparkPlan)
extends BroadcastExchangeLike
@@ -125,6 +126,11 @@ abstract class NativeBroadcastExchangeBase(mode: BroadcastMode, override val chi
val v = mode.transform(dataRows)
val dummyBroadcasted = new Broadcast[Any](-1) {
override protected def getValue(): Any = v
+
+ @sparkver("241kwaiae")
+ override def setValue(v: () => (Any, Long => Any)): Any =
+ throw new UnsupportedOperationException("dummyBroadcasted.setValue")
+
override protected def doUnpersist(blocking: Boolean): Unit = {
MethodUtils.invokeMethod(broadcast, true, "doUnpersist", Array(blocking))
}
diff --git a/spark-extension/src/main/scala/org/apache/spark/sql/execution/blaze/plan/NativeShuffledHashJoinBase.scala b/spark-extension/src/main/scala/org/apache/spark/sql/execution/blaze/plan/NativeShuffledHashJoinBase.scala
index fe42d51b..66ff94a3 100644
--- a/spark-extension/src/main/scala/org/apache/spark/sql/execution/blaze/plan/NativeShuffledHashJoinBase.scala
+++ b/spark-extension/src/main/scala/org/apache/spark/sql/execution/blaze/plan/NativeShuffledHashJoinBase.scala
@@ -110,7 +110,7 @@ abstract class NativeShuffledHashJoinBase(
partitions,
partitioner,
dependencies,
- leftRDD.isShuffleReadFull && rightRDD.isShuffleReadFull,
+ rddShuffleReadFull = false, // probed side will not be read if built side is empty
(partition, taskContext) => {
val leftPartition = leftRDD.partitions(partition.index)
val leftChild = leftRDD.nativePlan(leftPartition, taskContext)
diff --git a/spark-extension/src/main/scala/org/apache/spark/sql/hive/execution/blaze/plan/NativePaimonTableScanExec.scala b/spark-extension/src/main/scala/org/apache/spark/sql/hive/execution/blaze/plan/NativePaimonTableScanExec.scala
index cd67e622..f77efc36 100644
--- a/spark-extension/src/main/scala/org/apache/spark/sql/hive/execution/blaze/plan/NativePaimonTableScanExec.scala
+++ b/spark-extension/src/main/scala/org/apache/spark/sql/hive/execution/blaze/plan/NativePaimonTableScanExec.scala
@@ -16,6 +16,7 @@
package org.apache.spark.sql.hive.execution.blaze.plan
import java.util.UUID
+
import scala.collection.JavaConverters._
import scala.collection.Seq
@@ -42,6 +43,7 @@ import org.apache.spark.sql.hive.blaze.PaimonUtil
import org.apache.spark.sql.hive.execution.HiveTableScanExec
import org.apache.spark.sql.types.StructType
import org.blaze.{protobuf => pb}
+import org.blaze.sparkver
case class NativePaimonTableScanExec(basedHiveScan: HiveTableScanExec)
extends NativeHiveTableScanBase(basedHiveScan)
@@ -125,6 +127,12 @@ case class NativePaimonTableScanExec(basedHiveScan: HiveTableScanExec)
override val nodeName: String =
s"NativePaimonTableScan $tableName"
+ @sparkver("241kwaiae")
+ override def getFilePartitions(): Array[FilePartition] = {
+ throw new UnsupportedOperationException("spark-adaptive not supported")
+ }
+
+ @sparkver("3.0 / 3.1 / 3.2 / 3.3 / 3.4 / 3.5")
override def getFilePartitions(): Array[FilePartition] = {
val currentTimeMillis = System.currentTimeMillis()
val sparkSession = Shims.get.getSqlContext(basedHiveScan).sparkSession
(auron) 14/16: disable max/min + nested params conversion (https://team.corp.kuaishou.com/task/T10892667)
Posted by ri...@apache.org.
This is an automated email from the ASF dual-hosted git repository.
richox pushed a commit to branch dev-v6.0.0-parallel-scan-kdev-build
in repository https://gitbox.apache.org/repos/asf/auron.git
commit 08b6fa77309d2b42cc02c1968ba9dc8bb282abb6
Author: zhangli20 <ri...@gmail.com>
AuthorDate: Sat Mar 7 15:18:53 2026 +0800
disable max/min + nested params conversion (https://team.corp.kuaishou.com/task/T10892667)
---
.../org/apache/spark/sql/execution/blaze/plan/NativeAggExec.scala | 2 ++
.../main/scala/org/apache/spark/sql/blaze/NativeConverters.scala | 4 ++++
.../org/apache/spark/sql/execution/blaze/plan/NativeAggBase.scala | 4 +++-
.../apache/spark/sql/execution/blaze/plan/NativeWindowBase.scala | 6 +++++-
4 files changed, 14 insertions(+), 2 deletions(-)
diff --git a/spark-extension-shims-spark3/src/main/scala/org/apache/spark/sql/execution/blaze/plan/NativeAggExec.scala b/spark-extension-shims-spark3/src/main/scala/org/apache/spark/sql/execution/blaze/plan/NativeAggExec.scala
index 3dc3c4ee..6cfbbc36 100644
--- a/spark-extension-shims-spark3/src/main/scala/org/apache/spark/sql/execution/blaze/plan/NativeAggExec.scala
+++ b/spark-extension-shims-spark3/src/main/scala/org/apache/spark/sql/execution/blaze/plan/NativeAggExec.scala
@@ -47,6 +47,8 @@ case class NativeAggExec(
child)
with BaseAggregateExec {
+ override val output: Seq[Attribute] = outputAttributes
+
@sparkver("3.1 / 3.2 / 3.3 / 3.4 / 3.5")
override val requiredChildDistributionExpressions: Option[Seq[Expression]] =
theRequiredChildDistributionExpressions
diff --git a/spark-extension/src/main/scala/org/apache/spark/sql/blaze/NativeConverters.scala b/spark-extension/src/main/scala/org/apache/spark/sql/blaze/NativeConverters.scala
index f798dcee..2bcb756a 100644
--- a/spark-extension/src/main/scala/org/apache/spark/sql/blaze/NativeConverters.scala
+++ b/spark-extension/src/main/scala/org/apache/spark/sql/blaze/NativeConverters.scala
@@ -1057,9 +1057,13 @@ object NativeConverters extends Logging {
e.aggregateFunction match {
case e: Max =>
+ assert(!e.child.dataType.isInstanceOf[ArrayType] && !e.child.dataType.isInstanceOf[MapType] && !e.child.dataType.isInstanceOf[StructType],
+ s"Max on complex types (Array/Map/Struct) is not supported in native execution. Expression: ${e}")
aggBuilder.setAggFunction(pb.AggFunction.MAX)
aggBuilder.addChildren(convertExpr(e.child))
case e: Min =>
+ assert(!e.child.dataType.isInstanceOf[ArrayType] && !e.child.dataType.isInstanceOf[MapType] && !e.child.dataType.isInstanceOf[StructType],
+ s"Min on complex types (Array/Map/Struct) is not supported in native execution. Expression: ${e}")
aggBuilder.setAggFunction(pb.AggFunction.MIN)
aggBuilder.addChildren(convertExpr(e.child))
case e: Sum if e.dataType.isInstanceOf[AtomicType] =>
diff --git a/spark-extension/src/main/scala/org/apache/spark/sql/execution/blaze/plan/NativeAggBase.scala b/spark-extension/src/main/scala/org/apache/spark/sql/execution/blaze/plan/NativeAggBase.scala
index 2078117b..5824fac3 100644
--- a/spark-extension/src/main/scala/org/apache/spark/sql/execution/blaze/plan/NativeAggBase.scala
+++ b/spark-extension/src/main/scala/org/apache/spark/sql/execution/blaze/plan/NativeAggBase.scala
@@ -139,7 +139,7 @@ abstract class NativeAggBase(
nativeAggrs
nativeAggrModes
- override def output: Seq[Attribute] =
+ def outputAttributes: Seq[Attribute] =
if (nativeAggrModes.contains(pb.AggMode.FINAL)) {
groupingExpressions.map(_.toAttribute) ++ aggregateAttributes
} else {
@@ -149,6 +149,8 @@ abstract class NativeAggBase(
groupingExpressions.map(_.toAttribute) ++ aggBufferAttrs
}
+ override def output: Seq[Attribute] = outputAttributes
+
override def outputPartitioning: Partitioning =
child.outputPartitioning
diff --git a/spark-extension/src/main/scala/org/apache/spark/sql/execution/blaze/plan/NativeWindowBase.scala b/spark-extension/src/main/scala/org/apache/spark/sql/execution/blaze/plan/NativeWindowBase.scala
index 4ccf4112..4fd23354 100644
--- a/spark-extension/src/main/scala/org/apache/spark/sql/execution/blaze/plan/NativeWindowBase.scala
+++ b/spark-extension/src/main/scala/org/apache/spark/sql/execution/blaze/plan/NativeWindowBase.scala
@@ -17,7 +17,6 @@ package org.apache.spark.sql.execution.blaze.plan
import scala.collection.JavaConverters._
import scala.collection.immutable.SortedMap
-
import org.apache.spark.OneToOneDependency
import org.apache.spark.sql.blaze.MetricNode
import org.apache.spark.sql.blaze.NativeConverters
@@ -47,6 +46,7 @@ import org.apache.spark.sql.catalyst.expressions.aggregate.Count
import org.apache.spark.sql.catalyst.expressions.aggregate.Max
import org.apache.spark.sql.catalyst.expressions.aggregate.Min
import org.apache.spark.sql.catalyst.expressions.aggregate.Sum
+import org.apache.spark.sql.types.{ArrayType, MapType, StructType}
import org.blaze.protobuf.WindowGroupLimit
abstract class NativeWindowBase(
@@ -136,6 +136,8 @@ abstract class NativeWindowBase(
assert(
spec.frameSpecification == RowNumber().frame, // only supports RowFrame(Unbounded, CurrentRow)
s"window frame not supported: ${spec.frameSpecification}")
+ assert(!e.child.dataType.isInstanceOf[ArrayType] && !e.child.dataType.isInstanceOf[MapType] && !e.child.dataType.isInstanceOf[StructType],
+ s"Max on complex types (Array/Map/Struct) is not supported in native execution. Expression: ${e}")
windowExprBuilder.setFuncType(pb.WindowFunctionType.Agg)
windowExprBuilder.setAggFunc(pb.AggFunction.MAX)
windowExprBuilder.addChildren(NativeConverters.convertExpr(e.child))
@@ -144,6 +146,8 @@ abstract class NativeWindowBase(
assert(
spec.frameSpecification == RowNumber().frame, // only supports RowFrame(Unbounded, CurrentRow)
s"window frame not supported: ${spec.frameSpecification}")
+ assert(!e.child.dataType.isInstanceOf[ArrayType] && !e.child.dataType.isInstanceOf[MapType] && !e.child.dataType.isInstanceOf[StructType],
+ s"Min on complex types (Array/Map/Struct) is not supported in native execution. Expression: ${e}")
windowExprBuilder.setFuncType(pb.WindowFunctionType.Agg)
windowExprBuilder.setAggFunc(pb.AggFunction.MIN)
windowExprBuilder.addChildren(NativeConverters.convertExpr(e.child))
(auron) 11/16: fix possible deadlock in OnHeapSpillManager
Posted by ri...@apache.org.
This is an automated email from the ASF dual-hosted git repository.
richox pushed a commit to branch dev-v6.0.0-parallel-scan-kdev-build
in repository https://gitbox.apache.org/repos/asf/auron.git
commit 92216fc4441bfae24fe7c0f89761338f916d9ab4
Author: zhangli20 <ri...@gmail.com>
AuthorDate: Thu Dec 25 21:23:28 2025 +0800
fix possible deadlock in OnHeapSpillManager
---
.../spark/sql/blaze/memory/OnHeapSpill.scala | 38 ++++++++++++++++++----
.../sql/blaze/memory/OnHeapSpillManager.scala | 2 +-
2 files changed, 32 insertions(+), 8 deletions(-)
diff --git a/spark-extension/src/main/scala/org/apache/spark/sql/blaze/memory/OnHeapSpill.scala b/spark-extension/src/main/scala/org/apache/spark/sql/blaze/memory/OnHeapSpill.scala
index 137f67c4..3e6dac66 100644
--- a/spark-extension/src/main/scala/org/apache/spark/sql/blaze/memory/OnHeapSpill.scala
+++ b/spark-extension/src/main/scala/org/apache/spark/sql/blaze/memory/OnHeapSpill.scala
@@ -16,20 +16,32 @@
package org.apache.spark.sql.blaze.memory
import java.nio.ByteBuffer
-
import org.apache.spark.internal.Logging
+import org.apache.spark.memory.MemoryConsumer
+
+import java.util.concurrent.locks.ReentrantLock
case class OnHeapSpill(hsm: OnHeapSpillManager, id: Int) extends Logging {
private var spillBuf: SpillBuf = new MemBasedSpillBuf
+ private val lock = new ReentrantLock
def memUsed: Long = spillBuf.memUsed
def diskUsed: Long = spillBuf.diskUsed
def size: Long = spillBuf.size
def diskIOTime: Long = spillBuf.diskIOTime
+ private def withLock[T](f: => T): T = {
+ lock.lock()
+ try {
+ f
+ } finally {
+ lock.unlock()
+ }
+ }
+
def write(buf: ByteBuffer): Unit = {
var needSpill = false
- synchronized {
+ withLock {
spillBuf match {
case _: MemBasedSpillBuf =>
val acquiredMemory = hsm.acquireMemory(buf.capacity())
@@ -45,13 +57,13 @@ case class OnHeapSpill(hsm: OnHeapSpillManager, id: Int) extends Logging {
spillInternal()
}
- synchronized {
+ withLock {
spillBuf.write(buf)
}
}
def read(buf: ByteBuffer): Int = {
- synchronized {
+ withLock {
val oldMemUsed = memUsed
val startPosition = buf.position()
spillBuf.read(buf)
@@ -68,7 +80,7 @@ case class OnHeapSpill(hsm: OnHeapSpillManager, id: Int) extends Logging {
}
def release(): Unit = {
- synchronized {
+ withLock {
val oldMemUsed = memUsed
spillBuf = new ReleasedSpillBuf(spillBuf)
@@ -78,8 +90,20 @@ case class OnHeapSpill(hsm: OnHeapSpillManager, id: Int) extends Logging {
}
}
- def spill(): Long = {
- synchronized {
+ def spill(trigger: MemoryConsumer): Long = {
+ // this might have been locked if the spilling is triggered by OnHeapSpill.write
+ if (trigger == this.hsm) {
+ if (lock.tryLock()) {
+ try {
+ return spillInternal()
+ } finally {
+ lock.unlock()
+ }
+ }
+ return 0L
+ }
+
+ withLock {
spillInternal()
}
}
diff --git a/spark-extension/src/main/scala/org/apache/spark/sql/blaze/memory/OnHeapSpillManager.scala b/spark-extension/src/main/scala/org/apache/spark/sql/blaze/memory/OnHeapSpillManager.scala
index 5f89d710..81d105ab 100644
--- a/spark-extension/src/main/scala/org/apache/spark/sql/blaze/memory/OnHeapSpillManager.scala
+++ b/spark-extension/src/main/scala/org/apache/spark/sql/blaze/memory/OnHeapSpillManager.scala
@@ -158,7 +158,7 @@ class OnHeapSpillManager(taskContext: TaskContext)
val sortedSpills = spills.seq.sortBy(0 - _.map(_.memUsed).getOrElse(0L))
sortedSpills.foreach {
case Some(spill) if spill.memUsed > 0 =>
- totalFreed += spill.spill()
+ totalFreed += spill.spill(trigger)
if (totalFreed >= size) {
return totalFreed
}
(auron) 07/16: fix exception type
Posted by ri...@apache.org.
This is an automated email from the ASF dual-hosted git repository.
richox pushed a commit to branch dev-v6.0.0-parallel-scan-kdev-build
in repository https://gitbox.apache.org/repos/asf/auron.git
commit 5f43afda131b44ba00237575a675bb307f67d94b
Author: zhuangxian <zh...@kuaishou.com>
AuthorDate: Tue Dec 9 09:12:25 2025 +0000
fix exception type
KDev_MR_link：https://ksurl.cn/1d8JpqvU
---
.../execution/blaze/shuffle/celeborn/BlazeCelebornShuffleReader.scala | 3 +--
1 file changed, 1 insertion(+), 2 deletions(-)
diff --git a/spark-extension-shims-spark3/src/main/scala/org/apache/spark/sql/execution/blaze/shuffle/celeborn/BlazeCelebornShuffleReader.scala b/spark-extension-shims-spark3/src/main/scala/org/apache/spark/sql/execution/blaze/shuffle/celeborn/BlazeCelebornShuffleReader.scala
index 34f681cc..acefd46d 100644
--- a/spark-extension-shims-spark3/src/main/scala/org/apache/spark/sql/execution/blaze/shuffle/celeborn/BlazeCelebornShuffleReader.scala
+++ b/spark-extension-shims-spark3/src/main/scala/org/apache/spark/sql/execution/blaze/shuffle/celeborn/BlazeCelebornShuffleReader.scala
@@ -147,7 +147,6 @@ class BlazeCelebornShuffleReader[K, C](
reader.read().map { kv =>
val celebornInputStream = kv._2.asInstanceOf[CelebornInputStream]
-
// force disable decompression because compression is skipped in shuffle writer
try {
FieldUtils.writeField(
@@ -156,7 +155,7 @@ class BlazeCelebornShuffleReader[K, C](
Boolean.box(false).asInstanceOf[Object],
true)
} catch {
- case _: IllegalAccessException =>
+ case _: IllegalArgumentException =>
// ignore if field not found (likely in celeborn-0.6.0)
}
celebornInputStream
(auron) 05/16: Fix: keep backward compatibility with celeborn-054This reverts commit 437180b7506bb56da6ad0a77ba6f6e
Posted by ri...@apache.org.
This is an automated email from the ASF dual-hosted git repository.
richox pushed a commit to branch dev-v6.0.0-parallel-scan-kdev-build
in repository https://gitbox.apache.org/repos/asf/auron.git
commit 60328bc9e8899dd814714670d7b717f767362432
Author: zhangli20 <zh...@kuaishou.com>
AuthorDate: Mon Dec 8 09:19:57 2025 +0000
Fix: keep backward compatibility with celeborn-054This reverts commit 437180b7506bb56da6ad0a77ba6f6e
KDev_MR_link：https://ksurl.cn/aCmpUcJb
---
.../celeborn/BlazeCelebornShuffleReader.scala | 144 +++++++++++++++------
1 file changed, 104 insertions(+), 40 deletions(-)
diff --git a/spark-extension-shims-spark3/src/main/scala/org/apache/spark/sql/execution/blaze/shuffle/celeborn/BlazeCelebornShuffleReader.scala b/spark-extension-shims-spark3/src/main/scala/org/apache/spark/sql/execution/blaze/shuffle/celeborn/BlazeCelebornShuffleReader.scala
index 569061d3..34f681cc 100644
--- a/spark-extension-shims-spark3/src/main/scala/org/apache/spark/sql/execution/blaze/shuffle/celeborn/BlazeCelebornShuffleReader.scala
+++ b/spark-extension-shims-spark3/src/main/scala/org/apache/spark/sql/execution/blaze/shuffle/celeborn/BlazeCelebornShuffleReader.scala
@@ -25,6 +25,7 @@ import scala.reflect.ClassTag
import org.apache.celeborn.client.read.CelebornInputStream
import org.apache.celeborn.common.CelebornConf
+import org.apache.commons.lang3.reflect.FieldUtils
import org.apache.spark.TaskContext
import org.apache.spark.internal.Logging
import org.apache.spark.shuffle.ShuffleReadMetricsReporter
@@ -51,52 +52,115 @@ class BlazeCelebornShuffleReader[K, C](
with Logging {
override protected def readBlocks(): Iterator[InputStream] = {
- // force disable decompression because compression is skipped in shuffle writer
- val reader = new CelebornShuffleReader[K, C](
- handle,
- startPartition,
- endPartition,
- startMapIndex.getOrElse(0),
- endMapIndex.getOrElse(Int.MaxValue),
- context,
- conf,
- BlazeCelebornShuffleReader.createBypassingIncRecordsReadMetrics(metrics),
- shuffleIdTracker,
- false) {
-
- override def newSerializerInstance(dep: ShuffleDependency[K, _, C]): SerializerInstance = {
- new SerializerInstance {
- override def serialize[T: ClassTag](t: T): ByteBuffer =
- throw new UnsupportedOperationException(
- "BlazeCelebornShuffleReader.newSerializerInstance")
-
- override def deserialize[T: ClassTag](bytes: ByteBuffer): T =
- throw new UnsupportedOperationException(
- "BlazeCelebornShuffleReader.newSerializerInstance")
-
- override def deserialize[T: ClassTag](bytes: ByteBuffer, loader: ClassLoader): T =
- throw new UnsupportedOperationException(
- "BlazeCelebornShuffleReader.newSerializerInstance")
-
- override def serializeStream(s: OutputStream): SerializationStream =
- throw new UnsupportedOperationException(
- "BlazeCelebornShuffleReader.newSerializerInstance")
-
- override def deserializeStream(s: InputStream): DeserializationStream = {
- new DeserializationStream {
- override def asKeyValueIterator: Iterator[(Any, Any)] = Iterator.single((null, s))
-
- override def readObject[T: ClassTag](): T =
- throw new UnsupportedOperationException()
-
- override def close(): Unit = s.close()
+ val reader = {
+ try {
+ // for celeborn-060
+ new CelebornShuffleReader[K, C](
+ handle,
+ startPartition,
+ endPartition,
+ startMapIndex.getOrElse(0),
+ endMapIndex.getOrElse(Int.MaxValue),
+ context,
+ conf,
+ BlazeCelebornShuffleReader.createBypassingIncRecordsReadMetrics(metrics),
+ shuffleIdTracker, false) {
+
+ override def newSerializerInstance(dep: ShuffleDependency[K, _, C]): SerializerInstance = {
+ new SerializerInstance {
+ override def serialize[T: ClassTag](t: T): ByteBuffer =
+ throw new UnsupportedOperationException(
+ "BlazeCelebornShuffleReader.newSerializerInstance")
+
+ override def deserialize[T: ClassTag](bytes: ByteBuffer): T =
+ throw new UnsupportedOperationException(
+ "BlazeCelebornShuffleReader.newSerializerInstance")
+
+ override def deserialize[T: ClassTag](bytes: ByteBuffer, loader: ClassLoader): T =
+ throw new UnsupportedOperationException(
+ "BlazeCelebornShuffleReader.newSerializerInstance")
+
+ override def serializeStream(s: OutputStream): SerializationStream =
+ throw new UnsupportedOperationException(
+ "BlazeCelebornShuffleReader.newSerializerInstance")
+
+ override def deserializeStream(s: InputStream): DeserializationStream = {
+ new DeserializationStream {
+ override def asKeyValueIterator: Iterator[(Any, Any)] = Iterator.single((null, s))
+
+ override def readObject[T: ClassTag](): T =
+ throw new UnsupportedOperationException()
+
+ override def close(): Unit = s.close()
+ }
+ }
}
}
}
+ } catch {
+ case _: NoSuchMethodError =>
+ // for celeborn-054
+ new CelebornShuffleReader[K, C](
+ handle,
+ startPartition,
+ endPartition,
+ startMapIndex.getOrElse(0),
+ endMapIndex.getOrElse(Int.MaxValue),
+ context,
+ conf,
+ BlazeCelebornShuffleReader.createBypassingIncRecordsReadMetrics(metrics),
+ shuffleIdTracker) {
+
+ override def newSerializerInstance(dep: ShuffleDependency[K, _, C]): SerializerInstance = {
+ new SerializerInstance {
+ override def serialize[T: ClassTag](t: T): ByteBuffer =
+ throw new UnsupportedOperationException(
+ "BlazeCelebornShuffleReader.newSerializerInstance")
+
+ override def deserialize[T: ClassTag](bytes: ByteBuffer): T =
+ throw new UnsupportedOperationException(
+ "BlazeCelebornShuffleReader.newSerializerInstance")
+
+ override def deserialize[T: ClassTag](bytes: ByteBuffer, loader: ClassLoader): T =
+ throw new UnsupportedOperationException(
+ "BlazeCelebornShuffleReader.newSerializerInstance")
+
+ override def serializeStream(s: OutputStream): SerializationStream =
+ throw new UnsupportedOperationException(
+ "BlazeCelebornShuffleReader.newSerializerInstance")
+
+ override def deserializeStream(s: InputStream): DeserializationStream = {
+ new DeserializationStream {
+ override def asKeyValueIterator: Iterator[(Any, Any)] = Iterator.single((null, s))
+
+ override def readObject[T: ClassTag](): T =
+ throw new UnsupportedOperationException()
+
+ override def close(): Unit = s.close()
+ }
+ }
+ }
+ }
+ }
}
}
- reader.read().map { kv => kv._2.asInstanceOf[CelebornInputStream] }
+ reader.read().map { kv =>
+ val celebornInputStream = kv._2.asInstanceOf[CelebornInputStream]
+
+ // force disable decompression because compression is skipped in shuffle writer
+ try {
+ FieldUtils.writeField(
+ celebornInputStream,
+ "shuffleCompressionEnabled",
+ Boolean.box(false).asInstanceOf[Object],
+ true)
+ } catch {
+ case _: IllegalAccessException =>
+ // ignore if field not found (likely in celeborn-0.6.0)
+ }
+ celebornInputStream
+ }
}
}
(auron) 03/16: fix execute_projected_with_key_rows_output() (#1156)
Posted by ri...@apache.org.
This is an automated email from the ASF dual-hosted git repository.
richox pushed a commit to branch dev-v6.0.0-parallel-scan-kdev-build
in repository https://gitbox.apache.org/repos/asf/auron.git
commit 83e87e1e4f456d812e8686dac30a0a8c17b3c1b7
Author: Zhang Li <ri...@qq.com>
AuthorDate: Tue Aug 12 19:58:24 2025 +0800
fix execute_projected_with_key_rows_output() (#1156)
Co-authored-by: zhangli20 <zh...@zhangli20deMacBook-Pro.local>
---
.../datafusion-ext-plans/src/common/execution_context.rs | 4 ++--
.../datafusion-ext-plans/src/common/row_null_checker.rs | 9 ---------
2 files changed, 2 insertions(+), 11 deletions(-)
diff --git a/native-engine/datafusion-ext-plans/src/common/execution_context.rs b/native-engine/datafusion-ext-plans/src/common/execution_context.rs
index f21fc7d3..3da6526e 100644
--- a/native-engine/datafusion-ext-plans/src/common/execution_context.rs
+++ b/native-engine/datafusion-ext-plans/src/common/execution_context.rs
@@ -316,12 +316,12 @@ impl ExecutionContext {
);
}
- let output_schema = self.output_schema();
+ let input_schema = input.schema();
let key_converter = RowConverter::new(
keys.iter()
.map(|k| {
Ok(SortField::new_with_options(
- k.expr.data_type(&output_schema)?,
+ k.expr.data_type(&input_schema)?,
k.options,
))
})
diff --git a/native-engine/datafusion-ext-plans/src/common/row_null_checker.rs b/native-engine/datafusion-ext-plans/src/common/row_null_checker.rs
index 0a168635..dda5cc9a 100644
--- a/native-engine/datafusion-ext-plans/src/common/row_null_checker.rs
+++ b/native-engine/datafusion-ext-plans/src/common/row_null_checker.rs
@@ -438,15 +438,6 @@ mod tests {
use super::*;
- // Helper function to create a SortField for testing
- fn create_test_sort_field(data_type: DataType, nulls_first: bool) -> SortField {
- let sort_options = SortOptions {
- descending: false,
- nulls_first,
- };
- SortField::new_with_options(data_type, sort_options)
- }
-
#[test]
fn test_primitive_null_detection() {
// Create a SortField with Int32 data type
(auron) 02/16: fix execution error in non-native parquet sink tasks (#1123)
Posted by ri...@apache.org.
This is an automated email from the ASF dual-hosted git repository.
richox pushed a commit to branch dev-v6.0.0-parallel-scan-kdev-build
in repository https://gitbox.apache.org/repos/asf/auron.git
commit 0eaa812e107198689b79d8bd66cd30d5e0548b99
Author: Zhang Li <ri...@qq.com>
AuthorDate: Thu Aug 7 17:02:38 2025 +0800
fix execution error in non-native parquet sink tasks (#1123)
Co-authored-by: zhangli20 <zh...@kuaishou.com>
---
.../NativeParquetInsertIntoHiveTableExec.scala | 33 ++++++++++++++++++++--
.../NativeParquetInsertIntoHiveTableBase.scala | 1 +
.../blaze/plan/NativeParquetSinkBase.scala | 3 ++
3 files changed, 34 insertions(+), 3 deletions(-)
diff --git a/spark-extension-shims-spark3/src/main/scala/org/apache/spark/sql/execution/blaze/plan/NativeParquetInsertIntoHiveTableExec.scala b/spark-extension-shims-spark3/src/main/scala/org/apache/spark/sql/execution/blaze/plan/NativeParquetInsertIntoHiveTableExec.scala
index ef31c270..fccb22d9 100644
--- a/spark-extension-shims-spark3/src/main/scala/org/apache/spark/sql/execution/blaze/plan/NativeParquetInsertIntoHiveTableExec.scala
+++ b/spark-extension-shims-spark3/src/main/scala/org/apache/spark/sql/execution/blaze/plan/NativeParquetInsertIntoHiveTableExec.scala
@@ -113,9 +113,17 @@ case class NativeParquetInsertIntoHiveTableExec(
new BasicWriteJobStatsTracker(serializableHadoopConf, metrics) {
override def newTaskInstance(): WriteTaskStatsTracker = {
new BasicWriteTaskStatsTracker(serializableHadoopConf.value) {
- override def newRow(_filePath: String, _row: InternalRow): Unit = {}
+ override def newRow(filePath: String, row: InternalRow): Unit = {
+ if (!ParquetSinkTaskContext.get.isNative) {
+ return super.newRow(filePath, row)
+ }
+ }
override def closeFile(filePath: String): Unit = {
+ if (!ParquetSinkTaskContext.get.isNative) {
+ return super.closeFile(filePath)
+ }
+
val outputFileStat = ParquetSinkTaskContext.get.processedOutputFiles.remove()
for (_ <- 0L until outputFileStat.numRows) {
super.newRow(filePath, null)
@@ -147,12 +155,23 @@ case class NativeParquetInsertIntoHiveTableExec(
mutable.ArrayBuffer.empty
override def newPartition(partitionValues: InternalRow): Unit = {
+ if (!ParquetSinkTaskContext.get.isNative) {
+ return super.newPartition(partitionValues)
+ }
partitions.append(partitionValues)
}
- override def newRow(_row: InternalRow): Unit = {}
+ override def newRow(row: InternalRow): Unit = {
+ if (!ParquetSinkTaskContext.get.isNative) {
+ return super.newRow(row)
+ }
+ }
override def getFinalStats(): WriteTaskStats = {
+ if (!ParquetSinkTaskContext.get.isNative) {
+ return super.getFinalStats()
+ }
+
val outputFileStat = ParquetSinkTaskContext.get.processedOutputFiles.remove()
BasicWriteTaskStats(
partitions = partitions,
@@ -179,9 +198,17 @@ case class NativeParquetInsertIntoHiveTableExec(
new BasicWriteJobStatsTracker(serializableHadoopConf, metrics) {
override def newTaskInstance(): WriteTaskStatsTracker = {
new BasicWriteTaskStatsTracker(serializableHadoopConf.value) {
- override def newRow(_row: InternalRow): Unit = {}
+ override def newRow(row: InternalRow): Unit = {
+ if (!ParquetSinkTaskContext.get.isNative) {
+ return super.newRow(row)
+ }
+ }
override def getFinalStats(): WriteTaskStats = {
+ if (!ParquetSinkTaskContext.get.isNative) {
+ return super.getFinalStats()
+ }
+
val outputFileStat = ParquetSinkTaskContext.get.processedOutputFiles.remove()
BasicWriteTaskStats(
numPartitions = 1,
diff --git a/spark-extension/src/main/scala/org/apache/spark/sql/execution/blaze/plan/NativeParquetInsertIntoHiveTableBase.scala b/spark-extension/src/main/scala/org/apache/spark/sql/execution/blaze/plan/NativeParquetInsertIntoHiveTableBase.scala
index 8c4ee53d..cc8cd9bb 100644
--- a/spark-extension/src/main/scala/org/apache/spark/sql/execution/blaze/plan/NativeParquetInsertIntoHiveTableBase.scala
+++ b/spark-extension/src/main/scala/org/apache/spark/sql/execution/blaze/plan/NativeParquetInsertIntoHiveTableBase.scala
@@ -165,6 +165,7 @@ class BlazeMapredParquetOutputFormat
case class OutputFileStat(path: String, numRows: Long, numBytes: Long)
class ParquetSinkTaskContext {
+ var isNative: Boolean = false
val processingOutputFiles = new LinkedBlockingDeque[String]()
val processedOutputFiles = new util.ArrayDeque[OutputFileStat]()
}
diff --git a/spark-extension/src/main/scala/org/apache/spark/sql/execution/blaze/plan/NativeParquetSinkBase.scala b/spark-extension/src/main/scala/org/apache/spark/sql/execution/blaze/plan/NativeParquetSinkBase.scala
index f6b12886..1a405738 100644
--- a/spark-extension/src/main/scala/org/apache/spark/sql/execution/blaze/plan/NativeParquetSinkBase.scala
+++ b/spark-extension/src/main/scala/org/apache/spark/sql/execution/blaze/plan/NativeParquetSinkBase.scala
@@ -97,6 +97,9 @@ abstract class NativeParquetSinkBase(
inputRDD.isShuffleReadFull,
(partition, context) => {
+ // mark for native parquet sink
+ ParquetSinkTaskContext.get.isNative = true
+
// init hadoop fs
val resourceId = s"NativeParquetSinkExec:${UUID.randomUUID().toString}"
JniBridge.resourcesMap.put(
(auron) 13/16: disable string to decimal cast (https://team.corp.kuaishou.com/task/T10924138)
Posted by ri...@apache.org.
This is an automated email from the ASF dual-hosted git repository.
richox pushed a commit to branch dev-v6.0.0-parallel-scan-kdev-build
in repository https://gitbox.apache.org/repos/asf/auron.git
commit 186718c913945fb21119dd22e5cb8397948138e8
Author: zhangli20 <ri...@gmail.com>
AuthorDate: Sat Mar 7 15:14:11 2026 +0800
disable string to decimal cast (https://team.corp.kuaishou.com/task/T10924138)
---
pom.xml | 2 +-
.../src/main/scala/org/apache/spark/sql/blaze/NativeConverters.scala | 2 ++
2 files changed, 3 insertions(+), 1 deletion(-)
diff --git a/pom.xml b/pom.xml
index 1a186ce6..b034ec28 100644
--- a/pom.xml
+++ b/pom.xml
@@ -416,7 +416,7 @@
<id>spark-3.5</id>
<properties>
<shimName>spark-3.5</shimName>
- <pkgSuffix>-kwai-hotfix-rule-engine</pkgSuffix>
+ <pkgSuffix>-kwai-hotfix-decimal-cast</pkgSuffix>
<shimPkg>spark-extension-shims-spark3</shimPkg>
<javaVersion>1.8</javaVersion>
<scalaVersion>2.12</scalaVersion>
diff --git a/spark-extension/src/main/scala/org/apache/spark/sql/blaze/NativeConverters.scala b/spark-extension/src/main/scala/org/apache/spark/sql/blaze/NativeConverters.scala
index 4a57d474..f798dcee 100644
--- a/spark-extension/src/main/scala/org/apache/spark/sql/blaze/NativeConverters.scala
+++ b/spark-extension/src/main/scala/org/apache/spark/sql/blaze/NativeConverters.scala
@@ -437,6 +437,8 @@ object NativeConverters extends Logging {
case cast: Cast
if !Seq(cast.dataType, cast.child.dataType).exists(t =>
t.isInstanceOf[TimestampType] || t.isInstanceOf[DateType]) =>
+ assert(cast.child.dataType != StringType || !cast.dataType.isInstanceOf[DecimalType],
+ s"Cast from String to Decimal is not supported in native execution. Expression: ${cast}")
buildExprNode {
_.setTryCast(
pb.PhysicalTryCastNode
(auron) 12/16: Implement columnar aggregate buffers
Posted by ri...@apache.org.
This is an automated email from the ASF dual-hosted git repository.
richox pushed a commit to branch dev-v6.0.0-parallel-scan-kdev-build
in repository https://gitbox.apache.org/repos/asf/auron.git
commit 3fd3ba873251eed3125aa5569714075f8acfa92b
Author: zhangli20 <ri...@gmail.com>
AuthorDate: Wed Jan 21 17:48:44 2026 +0800
Implement columnar aggregate buffers
---
native-engine/datafusion-ext-plans/src/agg/acc.rs | 210 +++++++++------------
native-engine/datafusion-ext-plans/src/agg/agg.rs | 1 +
.../datafusion-ext-plans/src/agg/agg_ctx.rs | 84 +++++----
.../datafusion-ext-plans/src/agg/agg_table.rs | 12 +-
native-engine/datafusion-ext-plans/src/agg/avg.rs | 23 ++-
.../datafusion-ext-plans/src/agg/bloom_filter.rs | 43 +++--
.../src/agg/brickhouse/collect.rs | 4 +
.../src/agg/brickhouse/combine_unique.rs | 4 +
.../datafusion-ext-plans/src/agg/collect.rs | 40 +++-
.../datafusion-ext-plans/src/agg/count.rs | 143 +++-----------
.../datafusion-ext-plans/src/agg/first.rs | 36 ++--
.../src/agg/first_ignores_null.rs | 18 +-
.../datafusion-ext-plans/src/agg/maxmin.rs | 13 +-
.../src/agg/spark_udaf_wrapper.rs | 51 +++--
native-engine/datafusion-ext-plans/src/agg/sum.rs | 20 +-
.../sql/execution/blaze/plan/NativeAggExec.scala | 18 +-
.../spark/sql/auron/util/TaskContextHelper.scala | 135 +++++++++++++
.../apache/spark/sql/blaze/BlazeConverters.scala | 11 +-
.../sql/blaze/BlazeSparkSessionExtension.scala | 2 +-
.../org/apache/spark/sql/blaze/NativeHelper.scala | 3 +-
.../sql/execution/blaze/plan/NativeAggBase.scala | 53 +++---
21 files changed, 512 insertions(+), 412 deletions(-)
diff --git a/native-engine/datafusion-ext-plans/src/agg/acc.rs b/native-engine/datafusion-ext-plans/src/agg/acc.rs
index 7d8530a8..fe363e05 100644
--- a/native-engine/datafusion-ext-plans/src/agg/acc.rs
+++ b/native-engine/datafusion-ext-plans/src/agg/acc.rs
@@ -14,7 +14,7 @@
use std::{
any::Any,
- io::{Cursor, Read, Write},
+ io::{Read, Write},
sync::Arc,
};
@@ -23,7 +23,7 @@ use arrow::{
datatypes::{DataType, *},
};
use bitvec::{bitvec, vec::BitVec};
-use byteorder::{ReadBytesExt, WriteBytesExt};
+use byteorder::WriteBytesExt;
use datafusion::common::{Result, ScalarValue, utils::proxy::VecAllocExt};
use datafusion_ext_commons::{
SliceAsRawBytes, UninitializedInit, df_execution_err, downcast_any,
@@ -45,9 +45,9 @@ pub trait AccColumn: Send {
fn shrink_to_fit(&mut self);
fn num_records(&self) -> usize;
fn mem_used(&self) -> usize;
- fn freeze_to_rows(&self, idx: IdxSelection<'_>, array: &mut [Vec<u8>]) -> Result<()>;
- fn unfreeze_from_rows(&mut self, cursors: &mut [Cursor<&[u8]>]) -> Result<()>;
- fn spill(&self, idx: IdxSelection<'_>, w: &mut SpillCompressedWriter) -> Result<()>;
+ fn freeze_to_arrays(&mut self, idx: IdxSelection<'_>) -> Result<Vec<ArrayRef>>;
+ fn unfreeze_from_arrays(&mut self, arrays: &[ArrayRef]) -> Result<()>;
+ fn spill(&mut self, idx: IdxSelection<'_>, w: &mut SpillCompressedWriter) -> Result<()>;
fn unspill(&mut self, num_rows: usize, r: &mut SpillCompressedReader) -> Result<()>;
fn ensure_size(&mut self, idx: IdxSelection<'_>) {
@@ -139,8 +139,7 @@ impl AccBooleanColumn {
}
}
- pub fn to_array(&self, dt: &DataType, idx: IdxSelection<'_>) -> Result<ArrayRef> {
- assert!(dt == &DataType::Boolean);
+ pub fn to_array(&self, idx: IdxSelection<'_>) -> Result<ArrayRef> {
idx_with_iter!((idx @ idx) => {
Ok(Arc::new(BooleanArray::from_iter(
idx.map(|i| self.valids[i].then_some(self.values[i]))
@@ -176,38 +175,32 @@ impl AccColumn for AccBooleanColumn {
self.num_records() / 4 // 2 bits for each value
}
- fn freeze_to_rows(&self, idx: IdxSelection<'_>, array: &mut [Vec<u8>]) -> Result<()> {
- idx_with_iter!((idx @ idx) => {
- for (i, w) in idx.zip(array) {
- if self.valids[i] {
- w.write_u8(1 + self.values[i] as u8)?;
- } else {
- w.write_u8(0)?;
- }
- }
- });
- Ok(())
+ fn freeze_to_arrays(&mut self, idx: IdxSelection<'_>) -> Result<Vec<ArrayRef>> {
+ let array = self.to_array(idx)?;
+ Ok(vec![array])
}
- fn unfreeze_from_rows(&mut self, cursors: &mut [Cursor<&[u8]>]) -> Result<()> {
- self.resize(0);
+ fn unfreeze_from_arrays(&mut self, arrays: &[ArrayRef]) -> Result<()> {
+ let array = downcast_any!(&arrays[0], BooleanArray)?;
+ self.values.clear();
+ self.valids.clear();
- for cursor in cursors {
- match cursor.read_u8()? {
- 0 => {
- self.valids.push(false);
- self.values.push(false);
- }
- v => {
- self.valids.push(true);
- self.values.push(v - 1 != 0);
- }
+ // fill values
+ for i in 0..array.len() {
+ self.values.push(array.value(i));
+ }
+
+ // fill valids
+ if let Some(nb) = array.nulls() {
+ for bit in nb {
+ self.valids.push(bit);
}
}
+ self.valids.resize(self.values.len(), true);
Ok(())
}
- fn spill(&self, idx: IdxSelection<'_>, w: &mut SpillCompressedWriter) -> Result<()> {
+ fn spill(&mut self, idx: IdxSelection<'_>, w: &mut SpillCompressedWriter) -> Result<()> {
let mut buf = vec![];
idx_for! {
@@ -242,13 +235,15 @@ impl AccColumn for AccBooleanColumn {
pub struct AccPrimColumn<T: ArrowNativeType> {
values: Vec<T>,
valids: BitVec,
+ dt: DataType,
}
impl<T: ArrowNativeType> AccPrimColumn<T> {
- pub fn new(num_records: usize) -> Self {
+ pub fn new(num_records: usize, dt: DataType) -> Self {
Self {
values: vec![T::default(); num_records],
valids: bitvec![0; num_records],
+ dt,
}
}
@@ -331,39 +326,30 @@ impl<T: ArrowNativeType> AccColumn for AccPrimColumn<T> {
self.values.allocated_size() + (self.valids.capacity() + 7) / 8
}
- fn freeze_to_rows(&self, idx: IdxSelection<'_>, array: &mut [Vec<u8>]) -> Result<()> {
- idx_with_iter!((idx @ idx) => {
- for (i, w) in idx.zip(array) {
- if self.valids[i] {
- w.write_u8(1)?;
- w.write_all([self.values[i]].as_raw_bytes())?;
- } else {
- w.write_u8(0)?;
- }
- }
- });
- Ok(())
+ fn freeze_to_arrays(&mut self, idx: IdxSelection<'_>) -> Result<Vec<ArrayRef>> {
+ let array = self.to_array(&self.dt, idx)?;
+ Ok(vec![array])
}
- fn unfreeze_from_rows(&mut self, cursors: &mut [Cursor<&[u8]>]) -> Result<()> {
+ fn unfreeze_from_arrays(&mut self, arrays: &[ArrayRef]) -> Result<()> {
+ let array_data = arrays[0].to_data();
self.resize(0);
- let mut value_buf = [T::default()];
-
- for cursor in cursors {
- let valid = cursor.read_u8()?;
- if valid == 1 {
- cursor.read_exact(value_buf.as_raw_bytes_mut())?;
- self.values.push(value_buf[0]);
- self.valids.push(true);
- } else {
- self.values.push(T::default());
- self.valids.push(false);
+
+ // fill values
+ let buffer = array_data.buffer::<T>(0);
+ self.values = buffer[array_data.offset()..][..array_data.len()].to_vec();
+
+ // fill valids
+ if let Some(nb) = array_data.nulls() {
+ for bit in nb {
+ self.valids.push(bit);
}
}
+ self.valids.resize(self.values.len(), true);
Ok(())
}
- fn spill(&self, idx: IdxSelection<'_>, w: &mut SpillCompressedWriter) -> Result<()> {
+ fn spill(&mut self, idx: IdxSelection<'_>, w: &mut SpillCompressedWriter) -> Result<()> {
// write valids
let mut bits: BitVec<u8> = BitVec::with_capacity(idx.len());
idx_for! {
@@ -413,13 +399,15 @@ impl<T: ArrowNativeType> AccColumn for AccPrimColumn<T> {
pub struct AccBytesColumn {
items: Vec<Option<AccBytes>>,
heap_mem_used: usize,
+ dt: DataType,
}
impl AccBytesColumn {
- pub fn new(num_records: usize) -> Self {
+ pub fn new(num_records: usize, dt: DataType) -> Self {
Self {
items: vec![None; num_records],
heap_mem_used: 0,
+ dt,
}
}
@@ -438,13 +426,13 @@ impl AccBytesColumn {
self.heap_mem_used += self.item_heap_mem_used(idx);
}
- fn to_array(&self, dt: &DataType, idx: IdxSelection<'_>) -> Result<ArrayRef> {
+ fn to_array(&self, idx: IdxSelection<'_>) -> Result<ArrayRef> {
let binary;
idx_with_iter!((idx @ idx) => {
binary = BinaryArray::from_iter(idx.map(|i| self.items[i].as_ref()));
});
- match dt {
+ match &self.dt {
DataType::Utf8 => Ok(make_array(
binary
.to_data()
@@ -453,7 +441,7 @@ impl AccBytesColumn {
.build()?,
)),
DataType::Binary => Ok(Arc::new(binary)),
- _ => df_execution_err!("expected string or binary type, got {dt:?}"),
+ dt => df_execution_err!("expected string or binary type, got {dt:?}"),
}
}
@@ -534,25 +522,34 @@ impl AccColumn for AccBytesColumn {
self.heap_mem_used + self.items.allocated_size()
}
- fn freeze_to_rows(&self, idx: IdxSelection<'_>, array: &mut [Vec<u8>]) -> Result<()> {
- idx_with_iter!((idx @ idx) => {
- for (i, w) in idx.zip(array) {
- self.save_value(i, w)?;
- }
- });
- Ok(())
+ fn freeze_to_arrays(&mut self, idx: IdxSelection<'_>) -> Result<Vec<ArrayRef>> {
+ let array = self.to_array(idx)?;
+ Ok(vec![array])
}
- fn unfreeze_from_rows(&mut self, cursors: &mut [Cursor<&[u8]>]) -> Result<()> {
- self.items.resize(0, Default::default());
- for cursor in cursors {
- self.load_value(cursor)?;
+ fn unfreeze_from_arrays(&mut self, arrays: &[ArrayRef]) -> Result<()> {
+ let array_data = arrays[0].to_data();
+ self.items.clear();
+
+ // fill values
+ let data_buffer = array_data.buffer::<u8>(0);
+ let offset_buffer = array_data.buffer::<i32>(1);
+ for i in 0..array_data.len() {
+ if array_data.is_valid(i) {
+ let offset_begin = offset_buffer[array_data.offset() + i] as usize;
+ let offset_end = offset_buffer[array_data.offset() + i + 1] as usize;
+ self.items.push(Some(AccBytes::from_slice(
+ &data_buffer[offset_begin..offset_end],
+ )));
+ } else {
+ self.items.push(None);
+ }
}
self.refresh_heap_mem_used();
Ok(())
}
- fn spill(&self, idx: IdxSelection<'_>, w: &mut SpillCompressedWriter) -> Result<()> {
+ fn spill(&mut self, idx: IdxSelection<'_>, w: &mut SpillCompressedWriter) -> Result<()> {
idx_for! {
(idx in idx) => {
self.save_value(idx, w)?;
@@ -578,17 +575,17 @@ pub struct AccScalarValueColumn {
}
impl AccScalarValueColumn {
- pub fn new(dt: &DataType, num_rows: usize) -> Self {
- let null_value = ScalarValue::try_from(dt).expect("unsupported data type");
+ pub fn new(dt: DataType, num_rows: usize) -> Self {
+ let null_value = ScalarValue::try_from(&dt).expect("unsupported data type");
Self {
items: (0..num_rows).map(|_| null_value.clone()).collect(),
- dt: dt.clone(),
+ dt,
null_value,
heap_mem_used: 0,
}
}
- pub fn to_array(&mut self, _dt: &DataType, idx: IdxSelection<'_>) -> Result<ArrayRef> {
+ pub fn to_array(&mut self, idx: IdxSelection<'_>) -> Result<ArrayRef> {
idx_with_iter!((idx @ idx) => {
ScalarValue::iter_to_array(idx.map(|i| {
std::mem::replace(&mut self.items[i], self.null_value.clone())
@@ -644,38 +641,36 @@ impl AccColumn for AccScalarValueColumn {
self.heap_mem_used + self.items.allocated_size()
}
- fn freeze_to_rows(&self, idx: IdxSelection<'_>, array: &mut [Vec<u8>]) -> Result<()> {
- idx_with_iter!((idx @ idx) => {
- for (i, w) in idx.zip(array) {
- write_scalar(&self.items[i], true, w)?;
- }
- });
- Ok(())
+ fn freeze_to_arrays(&mut self, idx: IdxSelection<'_>) -> Result<Vec<ArrayRef>> {
+ let array = self.to_array(idx)?; // data type is not used
+ Ok(vec![array])
}
- fn unfreeze_from_rows(&mut self, cursors: &mut [Cursor<&[u8]>]) -> Result<()> {
- self.items.truncate(0);
+ fn unfreeze_from_arrays(&mut self, arrays: &[ArrayRef]) -> Result<()> {
+ let array = &arrays[0];
+
+ self.items.clear();
self.heap_mem_used = 0;
- for cursor in cursors {
- let scalar = read_scalar(cursor, &self.dt, true)?;
+ for i in 0..array.len() {
+ let scalar = ScalarValue::try_from_array(array, i)?;
self.heap_mem_used += scalar_value_heap_mem_size(&scalar);
- self.items.push(scalar);
}
Ok(())
}
- fn spill(&self, idx: IdxSelection<'_>, w: &mut SpillCompressedWriter) -> Result<()> {
+ fn spill(&mut self, idx: IdxSelection<'_>, w: &mut SpillCompressedWriter) -> Result<()> {
idx_for! {
(idx in idx) => {
- write_scalar(&self.items[idx], true, w)?;
+ let scalar = self.take_value(idx);
+ write_scalar(&scalar, true, w)?;
}
}
Ok(())
}
fn unspill(&mut self, num_rows: usize, r: &mut SpillCompressedReader) -> Result<()> {
- self.items.truncate(0);
+ self.items.clear();
self.heap_mem_used = 0;
for _ in 0..num_rows {
@@ -687,43 +682,18 @@ impl AccColumn for AccScalarValueColumn {
}
}
-pub fn create_acc_generic_column(dt: &DataType, num_rows: usize) -> AccColumnRef {
+pub fn create_acc_generic_column(dt: DataType, num_rows: usize) -> AccColumnRef {
macro_rules! primitive_helper {
($t:ty) => {
Box::new(AccPrimColumn::<<$t as ArrowPrimitiveType>::Native>::new(
- num_rows,
+ num_rows, dt,
))
};
}
downcast_primitive! {
dt => (primitive_helper),
DataType::Boolean => Box::new(AccBooleanColumn::new(num_rows)),
- DataType::Utf8 | DataType::Binary => Box::new(AccBytesColumn::new(num_rows)),
+ DataType::Utf8 | DataType::Binary => Box::new(AccBytesColumn::new(num_rows, dt)),
other => Box::new(AccScalarValueColumn::new(other, num_rows)),
}
}
-
-pub fn acc_generic_column_to_array(
- column: &mut AccColumnRef,
- dt: &DataType,
- idx: IdxSelection<'_>,
-) -> Result<ArrayRef> {
- macro_rules! primitive_helper {
- ($t:ty) => {
- downcast_any!(column, mut AccPrimColumn::<<$t as ArrowPrimitiveType>::Native>)?
- .to_array(dt, idx)
- };
- }
- downcast_primitive! {
- dt => (primitive_helper),
- DataType::Boolean => {
- downcast_any!(column, mut AccBooleanColumn)?.to_array(dt, idx)
- }
- DataType::Utf8 | DataType::Binary => {
- downcast_any!(column, mut AccBytesColumn)?.to_array(dt, idx)
- }
- _other => {
- downcast_any!(column, mut AccScalarValueColumn)?.to_array(dt, idx)
- }
- }
-}
diff --git a/native-engine/datafusion-ext-plans/src/agg/agg.rs b/native-engine/datafusion-ext-plans/src/agg/agg.rs
index 1e906143..82ab45d0 100644
--- a/native-engine/datafusion-ext-plans/src/agg/agg.rs
+++ b/native-engine/datafusion-ext-plans/src/agg/agg.rs
@@ -43,6 +43,7 @@ pub trait Agg: Send + Sync + Debug {
fn data_type(&self) -> &DataType;
fn nullable(&self) -> bool;
fn create_acc_column(&self, num_rows: usize) -> AccColumnRef;
+ fn acc_array_data_types(&self) -> &[DataType];
fn with_new_exprs(&self, exprs: Vec<PhysicalExprRef>) -> Result<Arc<dyn Agg>>;
fn prepare_partial_args(&self, partial_inputs: &[ArrayRef]) -> Result<Vec<ArrayRef>> {
diff --git a/native-engine/datafusion-ext-plans/src/agg/agg_ctx.rs b/native-engine/datafusion-ext-plans/src/agg/agg_ctx.rs
index 1efa7cd3..c34fe0eb 100644
--- a/native-engine/datafusion-ext-plans/src/agg/agg_ctx.rs
+++ b/native-engine/datafusion-ext-plans/src/agg/agg_ctx.rs
@@ -14,13 +14,12 @@
use std::{
fmt::{Debug, Formatter},
- io::Cursor,
sync::Arc,
};
use arrow::{
- array::{ArrayRef, BinaryArray, RecordBatchOptions},
- datatypes::{DataType, Field, Fields, Schema, SchemaRef},
+ array::{Array, ArrayRef, RecordBatchOptions},
+ datatypes::{Field, Fields, Schema, SchemaRef},
record_batch::RecordBatch,
row::{RowConverter, Rows, SortField},
};
@@ -28,20 +27,16 @@ use blaze_jni_bridge::{
conf,
conf::{BooleanConf, DoubleConf, IntConf},
};
-use datafusion::{
- common::{Result, cast::as_binary_array},
- physical_expr::PhysicalExprRef,
-};
+use datafusion::{common::Result, physical_expr::PhysicalExprRef};
use datafusion_ext_commons::{downcast_any, suggested_batch_mem_size};
use once_cell::sync::OnceCell;
use parking_lot::Mutex;
use crate::{
agg::{
- AGG_BUF_COLUMN_NAME, AggExecMode, AggExpr, AggMode, GroupingExpr,
+ AggExecMode, AggExpr, AggMode, GroupingExpr,
acc::AccTable,
agg::{Agg, IdxSelection},
- agg_hash_map::AggHashMapKey,
spark_udaf_wrapper::{AccUDAFBufferRowsColumn, SparkUDAFMemTracker, SparkUDAFWrapper},
},
common::{
@@ -62,6 +57,8 @@ pub struct AggContext {
pub grouping_row_converter: Arc<Mutex<RowConverter>>,
pub groupings: Vec<GroupingExpr>,
pub aggs: Vec<AggExpr>,
+ pub input_acc_arrays_len: usize,
+ pub output_acc_arrays_len: usize,
pub supports_partial_skipping: bool,
pub partial_skipping_ratio: f64,
pub partial_skipping_min_rows: usize,
@@ -136,7 +133,11 @@ impl AggContext {
));
}
} else {
- agg_fields.push(Field::new(AGG_BUF_COLUMN_NAME, DataType::Binary, false));
+ for agg in &aggs {
+ for dt in agg.agg.acc_array_data_types() {
+ agg_fields.push(Field::new("", dt.clone(), true));
+ }
+ }
}
let agg_schema = Arc::new(Schema::new(agg_fields));
let output_schema = Arc::new(Schema::new(
@@ -147,6 +148,17 @@ impl AggContext {
.concat(),
));
+ let input_acc_arrays_len = aggs
+ .iter()
+ .filter(|agg| agg.mode.is_partial_merge() || agg.mode.is_final())
+ .map(|agg| agg.agg.acc_array_data_types().len())
+ .sum();
+ let output_acc_arrays_len = aggs
+ .iter()
+ .filter(|agg| agg.mode.is_partial() || agg.mode.is_partial_merge())
+ .map(|agg| agg.agg.acc_array_data_types().len())
+ .sum();
+
let agg_exprs_flatten: Vec<PhysicalExprRef> = aggs
.iter()
.filter(|agg| agg.mode.is_partial())
@@ -194,6 +206,8 @@ impl AggContext {
grouping_row_converter,
groupings,
aggs,
+ input_acc_arrays_len,
+ output_acc_arrays_len,
agg_expr_evaluator,
supports_partial_skipping,
partial_skipping_ratio,
@@ -276,21 +290,15 @@ impl AggContext {
let mut merging_acc_table = self.create_acc_table(0);
if self.need_partial_merge {
- let partial_merged_array = as_binary_array(batch.columns().last().unwrap())?;
- let array = partial_merged_array
- .iter()
- .skip(batch_start_idx)
- .take(batch_end_idx - batch_start_idx)
- .map(|bytes| bytes.unwrap())
- .collect::<Vec<_>>();
- let mut cursors = array
- .iter()
- .map(|bytes| Cursor::new(bytes.as_bytes()))
- .collect::<Vec<_>>();
-
- for (agg_idx, _agg) in &self.need_partial_merge_aggs {
- let acc_col = &mut merging_acc_table.cols_mut()[*agg_idx];
- acc_col.unfreeze_from_rows(&mut cursors)?;
+ let mut acc_arrays_start = batch.num_columns() - self.input_acc_arrays_len;
+ for (agg_idx, agg) in &self.need_partial_merge_aggs {
+ let acc_arrays = &batch.columns()[acc_arrays_start..]
+ .iter()
+ .take(agg.acc_array_data_types().len())
+ .map(|array| array.slice(batch_start_idx, batch_end_idx - batch_start_idx))
+ .collect::<Vec<_>>();
+ acc_arrays_start += agg.acc_array_data_types().len();
+ merging_acc_table.cols_mut()[*agg_idx].unfreeze_from_arrays(&acc_arrays)?;
}
}
let batch_selection = IdxSelection::Range(0, batch_end_idx - batch_start_idx);
@@ -318,9 +326,7 @@ impl AggContext {
}
Ok(agg_columns)
} else {
- // output acc as a binary column
- let freezed = self.freeze_acc_table(acc_table, idx)?;
- Ok(vec![Arc::new(BinaryArray::from_iter_values(freezed))])
+ self.freeze_acc_table(acc_table, idx)
}
}
@@ -405,23 +411,23 @@ impl AggContext {
pub fn freeze_acc_table(
&self,
- acc_table: &AccTable,
+ acc_table: &mut AccTable,
acc_idx: IdxSelection,
- ) -> Result<Vec<Vec<u8>>> {
+ ) -> Result<Vec<ArrayRef>> {
let udaf_indices_cache = OnceCell::new();
- let mut vec = vec![vec![]; acc_idx.len()];
- for acc_col in acc_table.cols() {
+ let mut arrays = vec![];
+
+ for acc_col in acc_table.cols_mut() {
if let Ok(udaf_acc_col) = downcast_any!(acc_col, AccUDAFBufferRowsColumn) {
- udaf_acc_col.freeze_to_rows_with_indices_cache(
- acc_idx,
- &mut vec,
- &udaf_indices_cache,
- )?;
+ arrays.push(
+ udaf_acc_col
+ .freeze_to_array_with_indices_cache(acc_idx, &udaf_indices_cache)?,
+ );
} else {
- acc_col.freeze_to_rows(acc_idx, &mut vec)?;
+ arrays.extend(acc_col.freeze_to_arrays(acc_idx)?);
}
}
- Ok(vec)
+ Ok(arrays)
}
pub async fn process_partial_skipped(
diff --git a/native-engine/datafusion-ext-plans/src/agg/agg_table.rs b/native-engine/datafusion-ext-plans/src/agg/agg_table.rs
index 3eeaf651..861f94b1 100644
--- a/native-engine/datafusion-ext-plans/src/agg/agg_table.rs
+++ b/native-engine/datafusion-ext-plans/src/agg/agg_table.rs
@@ -543,7 +543,7 @@ impl HashingData {
// sort all records using radix sort on hashcodes of keys
let num_spill_buckets = self.agg_ctx.num_spill_buckets(self.mem_used());
let key_rows = self.map.into_keys();
- let acc_table = self.acc_table;
+ let mut acc_table = self.acc_table;
let mut entries = key_rows
.iter()
.enumerate()
@@ -567,7 +567,7 @@ impl HashingData {
write_spill_bucket(
&mut writer,
&self.agg_ctx,
- &acc_table,
+ &mut acc_table,
entries[offset..][..cur_bucket_count]
.iter()
.map(|&(_, record_idx)| &key_rows[record_idx as usize]),
@@ -673,7 +673,7 @@ impl MergingData {
entries.shrink_to_fit();
let key_rows = self.key_rows;
- let acc_table = self.acc_table;
+ let mut acc_table = self.acc_table;
let mut bucket_counts = vec![0; num_spill_buckets];
radix_sort_by_key(&mut entries, &mut bucket_counts, |(bucket_id, ..)| {
@@ -692,7 +692,7 @@ impl MergingData {
write_spill_bucket(
&mut writer,
&self.agg_ctx,
- &acc_table,
+ &mut acc_table,
entries[offset..][..cur_bucket_count].iter().map(
|&(_, batch_idx, row_idx, _)| {
key_rows[batch_idx as usize]
@@ -721,7 +721,7 @@ impl MergingData {
fn write_spill_bucket(
w: &mut SpillCompressedWriter,
agg_ctx: &AggContext,
- acc_table: &AccTable,
+ acc_table: &mut AccTable,
key_iter: impl Iterator<Item = impl AsRef<[u8]>>,
acc_idx_iter: impl Iterator<Item = usize>,
spill_idx: usize,
@@ -729,7 +729,7 @@ fn write_spill_bucket(
// write accs
let udaf_indices_cache = OnceCell::new();
let acc_indices: Vec<usize> = acc_idx_iter.collect();
- for col in acc_table.cols() {
+ for col in acc_table.cols_mut() {
if let Ok(udaf_col) = downcast_any!(col, AccUDAFBufferRowsColumn) {
udaf_col.spill_with_indices_cache(
IdxSelection::Indices(&acc_indices),
diff --git a/native-engine/datafusion-ext-plans/src/agg/avg.rs b/native-engine/datafusion-ext-plans/src/agg/avg.rs
index 1298737b..fb040468 100644
--- a/native-engine/datafusion-ext-plans/src/agg/avg.rs
+++ b/native-engine/datafusion-ext-plans/src/agg/avg.rs
@@ -15,7 +15,6 @@
use std::{
any::Any,
fmt::{Debug, Formatter},
- io::Cursor,
sync::Arc,
};
@@ -45,6 +44,7 @@ pub struct AggAvg {
data_type: DataType,
agg_sum: AggSum,
agg_count: AggCount,
+ acc_array_data_types: Vec<DataType>,
}
impl AggAvg {
@@ -52,6 +52,7 @@ impl AggAvg {
let agg_sum = AggSum::try_new(child.clone(), data_type.clone())?;
let agg_count = AggCount::try_new(vec![child.clone()], DataType::Int64)?;
Ok(Self {
+ acc_array_data_types: vec![data_type.clone(), DataType::Int64],
child,
data_type,
agg_sum,
@@ -105,6 +106,10 @@ impl Agg for AggAvg {
})
}
+ fn acc_array_data_types(&self) -> &[DataType] {
+ &self.acc_array_data_types
+ }
+
fn partial_update(
&self,
accs: &mut AccColumnRef,
@@ -204,19 +209,19 @@ impl AccColumn for AccAvgColumn {
self.sum.mem_used() + self.count.mem_used()
}
- fn freeze_to_rows(&self, idx: IdxSelection<'_>, array: &mut [Vec<u8>]) -> Result<()> {
- self.sum.freeze_to_rows(idx, array)?;
- self.count.freeze_to_rows(idx, array)?;
- Ok(())
+ fn freeze_to_arrays(&mut self, idx: IdxSelection<'_>) -> Result<Vec<ArrayRef>> {
+ let sum_array = self.sum.freeze_to_arrays(idx)?[0].clone();
+ let count_array = self.count.freeze_to_arrays(idx)?[0].clone();
+ Ok(vec![sum_array, count_array])
}
- fn unfreeze_from_rows(&mut self, cursors: &mut [Cursor<&[u8]>]) -> Result<()> {
- self.sum.unfreeze_from_rows(cursors)?;
- self.count.unfreeze_from_rows(cursors)?;
+ fn unfreeze_from_arrays(&mut self, arrays: &[ArrayRef]) -> Result<()> {
+ self.sum.unfreeze_from_arrays(&arrays[0..1])?;
+ self.count.unfreeze_from_arrays(&arrays[1..2])?;
Ok(())
}
- fn spill(&self, idx: IdxSelection<'_>, buf: &mut SpillCompressedWriter) -> Result<()> {
+ fn spill(&mut self, idx: IdxSelection<'_>, buf: &mut SpillCompressedWriter) -> Result<()> {
self.sum.spill(idx, buf)?;
self.count.spill(idx, buf)?;
Ok(())
diff --git a/native-engine/datafusion-ext-plans/src/agg/bloom_filter.rs b/native-engine/datafusion-ext-plans/src/agg/bloom_filter.rs
index 5e872970..f061ff60 100644
--- a/native-engine/datafusion-ext-plans/src/agg/bloom_filter.rs
+++ b/native-engine/datafusion-ext-plans/src/agg/bloom_filter.rs
@@ -20,7 +20,7 @@ use std::{
};
use arrow::{
- array::{ArrayRef, AsArray, BinaryBuilder},
+ array::{ArrayRef, AsArray, BinaryArray, BinaryBuilder},
datatypes::{DataType, Int64Type},
};
use byteorder::{ReadBytesExt, WriteBytesExt};
@@ -107,6 +107,10 @@ impl Agg for AggBloomFilter {
bloom_filters
}
+ fn acc_array_data_types(&self) -> &[DataType] {
+ &[DataType::Binary]
+ }
+
fn partial_update(
&self,
accs: &mut AccColumnRef,
@@ -244,36 +248,41 @@ impl AccColumn for AccBloomFilterColumn {
.sum()
}
- fn freeze_to_rows(&self, idx: IdxSelection<'_>, array: &mut [Vec<u8>]) -> Result<()> {
+ fn freeze_to_arrays(&mut self, idx: IdxSelection<'_>) -> Result<Vec<ArrayRef>> {
+ let mut array = vec![];
+
idx_for! {
(idx in idx) => {
- let w = &mut array[idx];
if let Some(bloom_filter) = &self.bloom_filters[idx] {
+ let mut w = vec![];
w.write_u8(1)?;
- bloom_filter.write_to(w)?;
+ bloom_filter.write_to(&mut w)?;
+ array.push(Some(w));
} else {
- w.write_u8(0)?;
+ array.push(None);
}
}
}
- Ok(())
+
+ let array = Arc::new(BinaryArray::from_iter(array));
+ Ok(vec![array])
}
- fn unfreeze_from_rows(&mut self, cursors: &mut [Cursor<&[u8]>]) -> Result<()> {
- assert_eq!(self.num_records(), 0, "expect empty AccColumn");
- for r in cursors {
- self.bloom_filters.push({
- if r.read_u8()? == 1 {
- Some(SparkBloomFilter::read_from(r)?)
- } else {
- None
- }
- });
+ fn unfreeze_from_arrays(&mut self, arrays: &[ArrayRef]) -> Result<()> {
+ let array = downcast_any!(arrays[0], BinaryArray)?;
+
+ for v in array {
+ if let Some(w) = v {
+ self.bloom_filters
+ .push(Some(SparkBloomFilter::read_from(&mut Cursor::new(w))?));
+ } else {
+ self.bloom_filters.push(None);
+ }
}
Ok(())
}
- fn spill(&self, idx: IdxSelection<'_>, w: &mut SpillCompressedWriter) -> Result<()> {
+ fn spill(&mut self, idx: IdxSelection<'_>, w: &mut SpillCompressedWriter) -> Result<()> {
idx_for! {
(idx in idx) => {
if let Some(bloom_filter) = &self.bloom_filters[idx] {
diff --git a/native-engine/datafusion-ext-plans/src/agg/brickhouse/collect.rs b/native-engine/datafusion-ext-plans/src/agg/brickhouse/collect.rs
index c9fd88d8..db9da487 100644
--- a/native-engine/datafusion-ext-plans/src/agg/brickhouse/collect.rs
+++ b/native-engine/datafusion-ext-plans/src/agg/brickhouse/collect.rs
@@ -80,6 +80,10 @@ impl Agg for AggCollect {
self.innert_collect_list.create_acc_column(num_rows)
}
+ fn acc_array_data_types(&self) -> &[DataType] {
+ self.innert_collect_list.acc_array_data_types()
+ }
+
fn partial_update(
&self,
accs: &mut AccColumnRef,
diff --git a/native-engine/datafusion-ext-plans/src/agg/brickhouse/combine_unique.rs b/native-engine/datafusion-ext-plans/src/agg/brickhouse/combine_unique.rs
index 67f3b6ca..ed993504 100644
--- a/native-engine/datafusion-ext-plans/src/agg/brickhouse/combine_unique.rs
+++ b/native-engine/datafusion-ext-plans/src/agg/brickhouse/combine_unique.rs
@@ -80,6 +80,10 @@ impl Agg for AggCombineUnique {
self.inner_collect_set.create_acc_column(num_rows)
}
+ fn acc_array_data_types(&self) -> &[DataType] {
+ self.inner_collect_set.acc_array_data_types()
+ }
+
fn partial_update(
&self,
accs: &mut AccColumnRef,
diff --git a/native-engine/datafusion-ext-plans/src/agg/collect.rs b/native-engine/datafusion-ext-plans/src/agg/collect.rs
index 88fc7270..341aea9d 100644
--- a/native-engine/datafusion-ext-plans/src/agg/collect.rs
+++ b/native-engine/datafusion-ext-plans/src/agg/collect.rs
@@ -115,6 +115,10 @@ impl<C: AccCollectionColumn> Agg for AggGenericCollect<C> {
col
}
+ fn acc_array_data_types(&self) -> &[DataType] {
+ &[DataType::Binary]
+ }
+
fn partial_update(
&self,
accs: &mut AccColumnRef,
@@ -291,15 +295,23 @@ impl AccColumn for AccSetColumn {
self.mem_used + self.set.capacity() * size_of::<AccSet>()
}
- fn freeze_to_rows(&self, idx: IdxSelection<'_>, array: &mut [Vec<u8>]) -> Result<()> {
- AccCollectionColumn::freeze_to_rows(self, idx, array)
+ fn freeze_to_arrays(&mut self, idx: IdxSelection<'_>) -> Result<Vec<ArrayRef>> {
+ let mut array = vec![vec![]; idx.len()];
+ AccCollectionColumn::freeze_to_rows(self, idx, &mut array)?;
+ Ok(vec![Arc::new(BinaryArray::from_iter_values(array))])
}
- fn unfreeze_from_rows(&mut self, cursors: &mut [Cursor<&[u8]>]) -> Result<()> {
- AccCollectionColumn::unfreeze_from_rows(self, cursors)
+ fn unfreeze_from_arrays(&mut self, arrays: &[ArrayRef]) -> Result<()> {
+ let array = downcast_any!(arrays[0], BinaryArray)?;
+ let mut cursors = vec![];
+
+ for i in 0..array.len() {
+ cursors.push(Cursor::new(array.value(i)));
+ }
+ AccCollectionColumn::unfreeze_from_rows(self, &mut cursors)
}
- fn spill(&self, idx: IdxSelection<'_>, w: &mut SpillCompressedWriter) -> Result<()> {
+ fn spill(&mut self, idx: IdxSelection<'_>, w: &mut SpillCompressedWriter) -> Result<()> {
idx_for! {
(idx in idx) => {
self.save_raw(idx, w)?;
@@ -403,15 +415,23 @@ impl AccColumn for AccListColumn {
self.mem_used + self.list.capacity() * size_of::<AccList>()
}
- fn freeze_to_rows(&self, idx: IdxSelection<'_>, array: &mut [Vec<u8>]) -> Result<()> {
- AccCollectionColumn::freeze_to_rows(self, idx, array)
+ fn freeze_to_arrays(&mut self, idx: IdxSelection<'_>) -> Result<Vec<ArrayRef>> {
+ let mut array = vec![vec![]; idx.len()];
+ AccCollectionColumn::freeze_to_rows(self, idx, &mut array)?;
+ Ok(vec![Arc::new(BinaryArray::from_iter_values(array))])
}
- fn unfreeze_from_rows(&mut self, cursors: &mut [Cursor<&[u8]>]) -> Result<()> {
- AccCollectionColumn::unfreeze_from_rows(self, cursors)
+ fn unfreeze_from_arrays(&mut self, arrays: &[ArrayRef]) -> Result<()> {
+ let array = downcast_any!(arrays[0], BinaryArray)?;
+ let mut cursors = vec![];
+
+ for i in 0..array.len() {
+ cursors.push(Cursor::new(array.value(i)));
+ }
+ AccCollectionColumn::unfreeze_from_rows(self, &mut cursors)
}
- fn spill(&self, idx: IdxSelection<'_>, w: &mut SpillCompressedWriter) -> Result<()> {
+ fn spill(&mut self, idx: IdxSelection<'_>, w: &mut SpillCompressedWriter) -> Result<()> {
idx_for! {
(idx in idx) => {
self.save_raw(idx, w)?;
diff --git a/native-engine/datafusion-ext-plans/src/agg/count.rs b/native-engine/datafusion-ext-plans/src/agg/count.rs
index 0ae34a76..61508c77 100644
--- a/native-engine/datafusion-ext-plans/src/agg/count.rs
+++ b/native-engine/datafusion-ext-plans/src/agg/count.rs
@@ -15,24 +15,19 @@
use std::{
any::Any,
fmt::{Debug, Formatter},
- io::Cursor,
sync::Arc,
};
use arrow::{array::*, datatypes::*};
use datafusion::{common::Result, physical_expr::PhysicalExprRef};
-use datafusion_ext_commons::{
- downcast_any,
- io::{read_len, write_len},
-};
+use datafusion_ext_commons::downcast_any;
use crate::{
agg::{
- acc::{AccColumn, AccColumnRef},
+ acc::{AccColumn, AccColumnRef, AccPrimColumn},
agg::{Agg, IdxSelection},
},
- idx_for, idx_for_zipped, idx_with_iter,
- memmgr::spill::{SpillCompressedReader, SpillCompressedWriter},
+ idx_for_zipped,
};
pub struct AggCount {
@@ -81,9 +76,11 @@ impl Agg for AggCount {
}
fn create_acc_column(&self, num_rows: usize) -> Box<dyn AccColumn> {
- Box::new(AccCountColumn {
- values: vec![0; num_rows],
- })
+ Box::new(AccPrimColumn::<i64>::new(num_rows, DataType::Int64))
+ }
+
+ fn acc_array_data_types(&self) -> &[DataType] {
+ &[DataType::Int64]
}
fn partial_update(
@@ -93,32 +90,15 @@ impl Agg for AggCount {
partial_args: &[ArrayRef],
partial_arg_idx: IdxSelection<'_>,
) -> Result<()> {
- let accs = downcast_any!(accs, mut AccCountColumn)?;
+ let accs = downcast_any!(accs, mut AccPrimColumn<i64>)?;
accs.ensure_size(acc_idx);
- if partial_args.is_empty() {
- idx_for_zipped! {
- ((acc_idx, _partial_arg_idx) in (acc_idx, partial_arg_idx)) => {
- if acc_idx >= accs.values.len() {
- accs.values.push(1);
- } else {
- accs.values[acc_idx] += 1;
- }
- }
- }
- } else {
- idx_for_zipped! {
- ((acc_idx, partial_arg_idx) in (acc_idx, partial_arg_idx)) => {
- let add = partial_args
- .iter()
- .all(|arg| arg.is_valid(partial_arg_idx)) as i64;
-
- if acc_idx >= accs.values.len() {
- accs.values.push(add);
- } else {
- accs.values[acc_idx] += add;
- }
- }
+ idx_for_zipped! {
+ ((acc_idx, partial_arg_idx) in (acc_idx, partial_arg_idx)) => {
+ let add = partial_args
+ .iter()
+ .all(|arg| arg.is_valid(partial_arg_idx)) as i64;
+ accs.set_value(acc_idx, Some(accs.value(acc_idx).unwrap_or(0) + add));
}
}
Ok(())
@@ -131,98 +111,25 @@ impl Agg for AggCount {
merging_accs: &mut AccColumnRef,
merging_acc_idx: IdxSelection<'_>,
) -> Result<()> {
- let accs = downcast_any!(accs, mut AccCountColumn)?;
- let merging_accs = downcast_any!(merging_accs, mut AccCountColumn)?;
+ let accs = downcast_any!(accs, mut AccPrimColumn<i64>)?;
+ let merging_accs = downcast_any!(merging_accs, mut AccPrimColumn<i64>)?;
accs.ensure_size(acc_idx);
idx_for_zipped! {
((acc_idx, merging_acc_idx) in (acc_idx, merging_acc_idx)) => {
- if acc_idx < accs.values.len() {
- accs.values[acc_idx] += merging_accs.values[merging_acc_idx];
- } else {
- accs.values.push(merging_accs.values[merging_acc_idx]);
- }
+ let v = match (accs.value(acc_idx), merging_accs.value(merging_acc_idx)) {
+ (Some(a), Some(b)) => Some(a + b),
+ (Some(a), _) => Some(a),
+ (_, Some(b)) => Some(b),
+ _ => Some(0),
+ };
+ accs.set_value(acc_idx, v);
}
}
Ok(())
}
fn final_merge(&self, accs: &mut AccColumnRef, acc_idx: IdxSelection<'_>) -> Result<ArrayRef> {
- let accs = downcast_any!(accs, mut AccCountColumn)?;
-
- idx_with_iter! {
- (acc_idx_iter @ acc_idx) => {
- Ok(Arc::new(Int64Array::from_iter_values(
- acc_idx_iter.map(|idx| accs.values[idx])
- )))
- }
- }
- }
-}
-
-pub struct AccCountColumn {
- pub values: Vec<i64>,
-}
-
-impl AccColumn for AccCountColumn {
- fn as_any(&self) -> &dyn Any {
- self
- }
-
- fn as_any_mut(&mut self) -> &mut dyn Any {
- self
- }
-
- fn resize(&mut self, num_accs: usize) {
- self.values.resize(num_accs, 0);
- }
-
- fn shrink_to_fit(&mut self) {
- self.values.shrink_to_fit();
- }
-
- fn num_records(&self) -> usize {
- self.values.len()
- }
-
- fn mem_used(&self) -> usize {
- self.values.capacity() * 2 * size_of::<i64>()
- }
-
- fn freeze_to_rows(&self, idx: IdxSelection<'_>, array: &mut [Vec<u8>]) -> Result<()> {
- let mut array_idx = 0;
-
- idx_for! {
- (idx in idx) => {
- write_len(self.values[idx] as usize, &mut array[array_idx])?;
- array_idx += 1;
- }
- }
- Ok(())
- }
-
- fn unfreeze_from_rows(&mut self, cursors: &mut [Cursor<&[u8]>]) -> Result<()> {
- assert_eq!(self.num_records(), 0, "expect empty AccColumn");
- for cursor in cursors {
- self.values.push(read_len(cursor)? as i64);
- }
- Ok(())
- }
-
- fn spill(&self, idx: IdxSelection<'_>, w: &mut SpillCompressedWriter) -> Result<()> {
- idx_for! {
- (idx in idx) => {
- write_len(self.values[idx] as usize, w)?;
- }
- }
- Ok(())
- }
-
- fn unspill(&mut self, num_rows: usize, r: &mut SpillCompressedReader) -> Result<()> {
- assert_eq!(self.num_records(), 0, "expect empty AccColumn");
- for _ in 0..num_rows {
- self.values.push(read_len(r)? as i64);
- }
- Ok(())
+ Ok(accs.freeze_to_arrays(acc_idx)?[0].clone())
}
}
diff --git a/native-engine/datafusion-ext-plans/src/agg/first.rs b/native-engine/datafusion-ext-plans/src/agg/first.rs
index 66b81eea..001e7ad8 100644
--- a/native-engine/datafusion-ext-plans/src/agg/first.rs
+++ b/native-engine/datafusion-ext-plans/src/agg/first.rs
@@ -15,7 +15,6 @@
use std::{
any::Any,
fmt::{Debug, Formatter},
- io::Cursor,
sync::Arc,
};
@@ -31,7 +30,7 @@ use crate::{
Agg,
acc::{
AccBooleanColumn, AccBytes, AccBytesColumn, AccColumn, AccColumnRef, AccPrimColumn,
- AccScalarValueColumn, acc_generic_column_to_array, create_acc_generic_column,
+ AccScalarValueColumn, create_acc_generic_column,
},
agg::IdxSelection,
},
@@ -42,11 +41,17 @@ use crate::{
pub struct AggFirst {
child: PhysicalExprRef,
data_type: DataType,
+ acc_array_data_types: Vec<DataType>,
}
impl AggFirst {
pub fn try_new(child: PhysicalExprRef, data_type: DataType) -> Result<Self> {
- Ok(Self { child, data_type })
+ let acc_array_data_types = vec![data_type.clone(), DataType::Boolean];
+ Ok(Self {
+ child,
+ data_type,
+ acc_array_data_types,
+ })
}
}
@@ -82,11 +87,15 @@ impl Agg for AggFirst {
fn create_acc_column(&self, num_rows: usize) -> AccColumnRef {
Box::new(AccFirstColumn {
- values: create_acc_generic_column(&self.data_type, num_rows),
+ values: create_acc_generic_column(self.data_type.clone(), num_rows),
flags: AccBooleanColumn::new(num_rows),
})
}
+ fn acc_array_data_types(&self) -> &[DataType] {
+ &self.acc_array_data_types
+ }
+
fn partial_update(
&self,
accs: &mut AccColumnRef,
@@ -266,8 +275,7 @@ impl Agg for AggFirst {
}
fn final_merge(&self, accs: &mut AccColumnRef, acc_idx: IdxSelection<'_>) -> Result<ArrayRef> {
- let accs = downcast_any!(accs, mut AccFirstColumn)?;
- acc_generic_column_to_array(&mut accs.values, &self.data_type, acc_idx)
+ Ok(accs.freeze_to_arrays(acc_idx)?[0].clone())
}
}
@@ -311,19 +319,19 @@ impl AccColumn for AccFirstColumn {
self.values.mem_used() + self.flags.mem_used()
}
- fn freeze_to_rows(&self, idx: IdxSelection<'_>, array: &mut [Vec<u8>]) -> Result<()> {
- self.values.freeze_to_rows(idx, array)?;
- self.flags.freeze_to_rows(idx, array)?;
- Ok(())
+ fn freeze_to_arrays(&mut self, idx: IdxSelection<'_>) -> Result<Vec<ArrayRef>> {
+ let value_array = self.values.freeze_to_arrays(idx)?[0].clone();
+ let flags_array = self.flags.freeze_to_arrays(idx)?[0].clone();
+ Ok(vec![value_array, flags_array])
}
- fn unfreeze_from_rows(&mut self, cursors: &mut [Cursor<&[u8]>]) -> Result<()> {
- self.values.unfreeze_from_rows(cursors)?;
- self.flags.unfreeze_from_rows(cursors)?;
+ fn unfreeze_from_arrays(&mut self, arrays: &[ArrayRef]) -> Result<()> {
+ self.values.unfreeze_from_arrays(&arrays[0..1])?;
+ self.flags.unfreeze_from_arrays(&arrays[1..2])?;
Ok(())
}
- fn spill(&self, idx: IdxSelection<'_>, w: &mut SpillCompressedWriter) -> Result<()> {
+ fn spill(&mut self, idx: IdxSelection<'_>, w: &mut SpillCompressedWriter) -> Result<()> {
self.values.spill(idx, w)?;
self.flags.spill(idx, w)?;
Ok(())
diff --git a/native-engine/datafusion-ext-plans/src/agg/first_ignores_null.rs b/native-engine/datafusion-ext-plans/src/agg/first_ignores_null.rs
index cb1d5e71..e6966e21 100644
--- a/native-engine/datafusion-ext-plans/src/agg/first_ignores_null.rs
+++ b/native-engine/datafusion-ext-plans/src/agg/first_ignores_null.rs
@@ -27,7 +27,7 @@ use crate::{
Agg,
acc::{
AccBooleanColumn, AccBytes, AccBytesColumn, AccColumnRef, AccPrimColumn,
- AccScalarValueColumn, acc_generic_column_to_array, create_acc_generic_column,
+ AccScalarValueColumn, create_acc_generic_column,
},
agg::IdxSelection,
},
@@ -37,11 +37,17 @@ use crate::{
pub struct AggFirstIgnoresNull {
child: PhysicalExprRef,
data_type: DataType,
+ acc_array_data_types: Vec<DataType>,
}
impl AggFirstIgnoresNull {
pub fn try_new(child: PhysicalExprRef, data_type: DataType) -> Result<Self> {
- Ok(Self { child, data_type })
+ let acc_array_data_types = vec![data_type.clone()];
+ Ok(Self {
+ child,
+ data_type,
+ acc_array_data_types,
+ })
}
}
@@ -76,7 +82,11 @@ impl Agg for AggFirstIgnoresNull {
}
fn create_acc_column(&self, num_rows: usize) -> AccColumnRef {
- create_acc_generic_column(&self.data_type, num_rows)
+ create_acc_generic_column(self.data_type.clone(), num_rows)
+ }
+
+ fn acc_array_data_types(&self) -> &[DataType] {
+ &self.acc_array_data_types
}
fn partial_update(
@@ -216,6 +226,6 @@ impl Agg for AggFirstIgnoresNull {
}
fn final_merge(&self, accs: &mut AccColumnRef, acc_idx: IdxSelection<'_>) -> Result<ArrayRef> {
- acc_generic_column_to_array(accs, &self.data_type, acc_idx)
+ Ok(accs.freeze_to_arrays(acc_idx)?[0].clone())
}
}
diff --git a/native-engine/datafusion-ext-plans/src/agg/maxmin.rs b/native-engine/datafusion-ext-plans/src/agg/maxmin.rs
index efe6c39a..ae9c5ab0 100644
--- a/native-engine/datafusion-ext-plans/src/agg/maxmin.rs
+++ b/native-engine/datafusion-ext-plans/src/agg/maxmin.rs
@@ -29,7 +29,7 @@ use crate::{
Agg,
acc::{
AccBooleanColumn, AccBytes, AccBytesColumn, AccColumn, AccColumnRef, AccPrimColumn,
- AccScalarValueColumn, acc_generic_column_to_array, create_acc_generic_column,
+ AccScalarValueColumn, create_acc_generic_column,
},
agg::IdxSelection,
},
@@ -42,14 +42,17 @@ pub type AggMin = AggMaxMin<AggMinParams>;
pub struct AggMaxMin<P: AggMaxMinParams> {
child: PhysicalExprRef,
data_type: DataType,
+ acc_array_data_types: Vec<DataType>,
_phantom: PhantomData<P>,
}
impl<P: AggMaxMinParams> AggMaxMin<P> {
pub fn try_new(child: PhysicalExprRef, data_type: DataType) -> Result<Self> {
+ let acc_array_data_types = vec![data_type.clone()];
Ok(Self {
child,
data_type,
+ acc_array_data_types,
_phantom: Default::default(),
})
}
@@ -86,7 +89,11 @@ impl<P: AggMaxMinParams> Agg for AggMaxMin<P> {
}
fn create_acc_column(&self, num_rows: usize) -> AccColumnRef {
- create_acc_generic_column(&self.data_type, num_rows)
+ create_acc_generic_column(self.data_type.clone(), num_rows)
+ }
+
+ fn acc_array_data_types(&self) -> &[DataType] {
+ &self.acc_array_data_types
}
fn partial_update(
@@ -285,7 +292,7 @@ impl<P: AggMaxMinParams> Agg for AggMaxMin<P> {
}
fn final_merge(&self, accs: &mut AccColumnRef, acc_idx: IdxSelection<'_>) -> Result<ArrayRef> {
- acc_generic_column_to_array(accs, &self.data_type, acc_idx)
+ Ok(accs.freeze_to_arrays(acc_idx)?[0].clone())
}
}
diff --git a/native-engine/datafusion-ext-plans/src/agg/spark_udaf_wrapper.rs b/native-engine/datafusion-ext-plans/src/agg/spark_udaf_wrapper.rs
index 554cbff2..c891285f 100644
--- a/native-engine/datafusion-ext-plans/src/agg/spark_udaf_wrapper.rs
+++ b/native-engine/datafusion-ext-plans/src/agg/spark_udaf_wrapper.rs
@@ -20,7 +20,10 @@ use std::{
};
use arrow::{
- array::{Array, ArrayRef, StructArray, as_struct_array, make_array},
+ array::{
+ Array, ArrayAccessor, ArrayRef, BinaryArray, BinaryBuilder, StructArray, as_struct_array,
+ make_array,
+ },
datatypes::{DataType, Field, Schema, SchemaRef},
ffi::{FFI_ArrowArray, FFI_ArrowSchema, from_ffi},
record_batch::{RecordBatch, RecordBatchOptions},
@@ -237,6 +240,10 @@ impl Agg for SparkUDAFWrapper {
Box::new(AccUDAFBufferRowsColumn { obj, jcontext })
}
+ fn acc_array_data_types(&self) -> &[DataType] {
+ &[DataType::Binary]
+ }
+
fn with_new_exprs(&self, _exprs: Vec<PhysicalExprRef>) -> Result<Arc<dyn Agg>> {
Ok(Arc::new(Self::try_new(
self.serialized.clone(),
@@ -288,12 +295,11 @@ pub struct AccUDAFBufferRowsColumn {
}
impl AccUDAFBufferRowsColumn {
- pub fn freeze_to_rows_with_indices_cache(
+ pub fn freeze_to_array_with_indices_cache(
&self,
idx: IdxSelection<'_>,
- array: &mut [Vec<u8>],
cache: &OnceCell<LocalRef>,
- ) -> Result<()> {
+ ) -> Result<ArrayRef> {
let idx_array =
cache.get_or_try_init(move || jni_new_prim_array!(int, &idx.to_int32_vec()[..]))?;
let serialized = jni_call!(
@@ -306,15 +312,18 @@ impl AccUDAFBufferRowsColumn {
jni_get_byte_array_region!(serialized.as_obj(), 0, &mut serialized_bytes[..])?;
// UnsafeRow is serialized with big-endian i32 length prefix
- let mut cursor = Cursor::new(&serialized_bytes);
- for i in 0..array.len() {
- let mut bytes_len_buf = [0; 4];
- cursor.read_exact(&mut bytes_len_buf)?;
+ let mut serialized_pos = 0;
+ let mut binary_builder = BinaryBuilder::with_capacity(idx.len(), 0);
+ for i in 0..idx.len() {
+ let mut bytes_len_buf = [0u8; 4];
+ bytes_len_buf.copy_from_slice(&serialized_bytes[serialized_pos..][..4]);
let bytes_len = i32::from_be_bytes(bytes_len_buf) as usize;
- write_len(bytes_len, &mut array[i])?;
- std::io::copy(&mut (&mut cursor).take(bytes_len as u64), &mut array[i])?;
+ serialized_pos += 4;
+
+ binary_builder.append_value(&serialized_bytes[serialized_pos..][..bytes_len]);
+ serialized_pos += bytes_len;
}
- Ok(())
+ Ok(Arc::new(binary_builder.finish()))
}
pub fn spill_with_indices_cache(
@@ -388,15 +397,23 @@ impl AccColumn for AccUDAFBufferRowsColumn {
0 // memory is managed in jvm side
}
- fn freeze_to_rows(&self, idx: IdxSelection<'_>, array: &mut [Vec<u8>]) -> Result<()> {
- self.freeze_to_rows_with_indices_cache(idx, array, &OnceCell::new())
+ fn freeze_to_arrays(&mut self, idx: IdxSelection<'_>) -> Result<Vec<ArrayRef>> {
+ let array = self.freeze_to_array_with_indices_cache(idx, &OnceCell::new())?;
+ Ok(vec![array])
}
- fn unfreeze_from_rows(&mut self, cursors: &mut [Cursor<&[u8]>]) -> Result<()> {
+ fn unfreeze_from_arrays(&mut self, arrays: &[ArrayRef]) -> Result<()> {
assert_eq!(self.num_records(), 0, "expect empty AccColumn");
+ let array = downcast_any!(arrays[0], BinaryArray)?;
+
+ let mut cursors = vec![];
+ for i in 0..array.len() {
+ cursors.push(Cursor::new(array.value(i)));
+ }
+
let mut data = vec![];
- for cursor in cursors.iter_mut() {
- let bytes_len = read_len(cursor)?;
+ for (i, cursor) in cursors.iter_mut().enumerate() {
+ let bytes_len = array.value(i).len();
data.write_all((bytes_len as i32).to_be_bytes().as_ref())?;
std::io::copy(&mut cursor.take(bytes_len as u64), &mut data)?;
}
@@ -413,7 +430,7 @@ impl AccColumn for AccUDAFBufferRowsColumn {
Ok(())
}
- fn spill(&self, _idx: IdxSelection<'_>, _buf: &mut SpillCompressedWriter) -> Result<()> {
+ fn spill(&mut self, _idx: IdxSelection<'_>, _buf: &mut SpillCompressedWriter) -> Result<()> {
unimplemented!("should call spill_with_indices_cache instead")
}
diff --git a/native-engine/datafusion-ext-plans/src/agg/sum.rs b/native-engine/datafusion-ext-plans/src/agg/sum.rs
index 7e1470d4..f373eb4a 100644
--- a/native-engine/datafusion-ext-plans/src/agg/sum.rs
+++ b/native-engine/datafusion-ext-plans/src/agg/sum.rs
@@ -25,9 +25,7 @@ use datafusion_ext_commons::{df_unimplemented_err, downcast_any};
use crate::{
agg::{
Agg,
- acc::{
- AccColumnRef, AccPrimColumn, acc_generic_column_to_array, create_acc_generic_column,
- },
+ acc::{AccColumnRef, AccPrimColumn, create_acc_generic_column},
agg::IdxSelection,
},
idx_for_zipped,
@@ -36,11 +34,17 @@ use crate::{
pub struct AggSum {
child: PhysicalExprRef,
data_type: DataType,
+ acc_array_data_types: Vec<DataType>,
}
impl AggSum {
pub fn try_new(child: PhysicalExprRef, data_type: DataType) -> Result<Self> {
- Ok(Self { child, data_type })
+ let acc_array_data_types = vec![data_type.clone()];
+ Ok(Self {
+ child,
+ data_type,
+ acc_array_data_types,
+ })
}
}
@@ -83,7 +87,11 @@ impl Agg for AggSum {
}
fn create_acc_column(&self, num_rows: usize) -> AccColumnRef {
- create_acc_generic_column(&self.data_type, num_rows)
+ create_acc_generic_column(self.data_type.clone(), num_rows)
+ }
+
+ fn acc_array_data_types(&self) -> &[DataType] {
+ &self.acc_array_data_types
}
fn partial_update(
@@ -144,6 +152,6 @@ impl Agg for AggSum {
}
fn final_merge(&self, accs: &mut AccColumnRef, acc_idx: IdxSelection<'_>) -> Result<ArrayRef> {
- acc_generic_column_to_array(accs, &self.data_type, acc_idx)
+ Ok(accs.freeze_to_arrays(acc_idx)?[0].clone())
}
}
diff --git a/spark-extension-shims-spark3/src/main/scala/org/apache/spark/sql/execution/blaze/plan/NativeAggExec.scala b/spark-extension-shims-spark3/src/main/scala/org/apache/spark/sql/execution/blaze/plan/NativeAggExec.scala
index 51625490..3dc3c4ee 100644
--- a/spark-extension-shims-spark3/src/main/scala/org/apache/spark/sql/execution/blaze/plan/NativeAggExec.scala
+++ b/spark-extension-shims-spark3/src/main/scala/org/apache/spark/sql/execution/blaze/plan/NativeAggExec.scala
@@ -15,17 +15,18 @@
*/
package org.apache.spark.sql.execution.blaze.plan
-import org.apache.spark.sql.catalyst.expressions.Attribute
import org.apache.spark.sql.catalyst.expressions.aggregate.AggregateExpression
+import org.apache.spark.sql.catalyst.expressions.aggregate.Final
+import org.apache.spark.sql.catalyst.expressions.Attribute
import org.apache.spark.sql.catalyst.expressions.AttributeReference
import org.apache.spark.sql.catalyst.expressions.Expression
-import org.apache.spark.sql.catalyst.expressions.ExprId
import org.apache.spark.sql.catalyst.expressions.NamedExpression
-import org.apache.spark.sql.catalyst.expressions.aggregate.Final
-import org.apache.spark.sql.execution.SparkPlan
import org.apache.spark.sql.execution.aggregate.BaseAggregateExec
+import org.apache.spark.sql.execution.auron.plan.NativeAggBase.AggExecMode
import org.apache.spark.sql.execution.blaze.plan.NativeAggBase.AggExecMode
+import org.apache.spark.sql.execution.SparkPlan
import org.apache.spark.sql.types.BinaryType
+
import org.blaze.sparkver
case class NativeAggExec(
@@ -53,15 +54,6 @@ case class NativeAggExec(
@sparkver("3.3 / 3.4 / 3.5")
override val initialInputBufferOffset: Int = theInitialInputBufferOffset
- override def output: Seq[Attribute] =
- if (aggregateExpressions.map(_.mode).contains(Final)) {
- groupingExpressions.map(_.toAttribute) ++ aggregateAttributes
- } else {
- groupingExpressions.map(_.toAttribute) :+
- AttributeReference(NativeAggBase.AGG_BUF_COLUMN_NAME, BinaryType, nullable = false)(
- ExprId.apply(NativeAggBase.AGG_BUF_COLUMN_EXPR_ID))
- }
-
@sparkver("3.2 / 3.3 / 3.4 / 3.5")
override def isStreaming: Boolean = false
diff --git a/spark-extension/src/main/scala/org/apache/spark/sql/auron/util/TaskContextHelper.scala b/spark-extension/src/main/scala/org/apache/spark/sql/auron/util/TaskContextHelper.scala
new file mode 100644
index 00000000..eecec5ef
--- /dev/null
+++ b/spark-extension/src/main/scala/org/apache/spark/sql/auron/util/TaskContextHelper.scala
@@ -0,0 +1,135 @@
+/*
+ * 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.auron.util
+
+import scala.util.control.NonFatal
+
+import org.apache.spark.{SparkEnv, TaskContext}
+import org.apache.spark.deploy.SparkHadoopUtil
+import org.apache.spark.internal.Logging
+import org.apache.spark.internal.config.APP_CALLER_CONTEXT
+import org.apache.spark.util.Utils
+
+object TaskContextHelper extends Logging {
+
+ private val callerContextSupported: Boolean = {
+ SparkHadoopUtil.get.conf.getBoolean("hadoop.caller.context.enabled", false) && {
+ try {
+ Utils.classForName("org.apache.hadoop.ipc.CallerContext")
+ Utils.classForName("org.apache.hadoop.ipc.CallerContext$Builder")
+ true
+ } catch {
+ case _: ClassNotFoundException =>
+ false
+ case NonFatal(e) =>
+ logWarning("Fail to load the CallerContext class", e)
+ false
+ }
+ }
+ }
+
+ def setNativeThreadName(): Unit = {
+ val context: TaskContext = TaskContext.get()
+ val thread = Thread.currentThread()
+ val threadName = if (context != null) {
+ s"auron native task ${context.partitionId()}.${context.attemptNumber()} in stage ${context
+ .stageId()}.${context.stageAttemptNumber()} (TID ${context.taskAttemptId()})"
+ } else {
+ "auron native task " + thread.getName
+ }
+ thread.setName(threadName)
+ }
+
+ def setHDFSCallerContext(): Unit = {
+ if (!callerContextSupported) {
+ return
+ }
+ val context: TaskContext = TaskContext.get()
+ if (context != null) {
+ val conf = SparkEnv.get.conf
+ val appId = conf.get("spark.app.id", "")
+ val appAttemptId = conf.get("spark.app.attempt.id", "")
+ // Spark executor cannot get the jobId from TaskContext, so we set a default value -1 here.
+ val jobId = -1
+ new CallerContextHelper(
+ "TASK",
+ conf.get(APP_CALLER_CONTEXT),
+ Option(appId),
+ if (appAttemptId == "") None else Option(appAttemptId),
+ Option(jobId),
+ Option(context.stageId()),
+ Option(context.stageAttemptNumber()),
+ Option(context.taskAttemptId()),
+ Option(context.attemptNumber())).setCurrentContext()
+ }
+ }
+
+ /**
+ * Copied from Apache Spark org.apache.spark.util.CallerContext
+ */
+ private class CallerContextHelper(
+ from: String,
+ upstreamCallerContext: Option[String] = None,
+ appId: Option[String] = None,
+ appAttemptId: Option[String] = None,
+ jobId: Option[Int] = None,
+ stageId: Option[Int] = None,
+ stageAttemptId: Option[Int] = None,
+ taskId: Option[Long] = None,
+ taskAttemptNumber: Option[Int] = None)
+ extends Logging {
+
+ private val context = prepareContext(
+ "SPARK_" +
+ from +
+ appId.map("_" + _).getOrElse("") +
+ appAttemptId.map("_" + _).getOrElse("") +
+ jobId.map("_JId_" + _).getOrElse("") +
+ stageId.map("_SId_" + _).getOrElse("") +
+ stageAttemptId.map("_" + _).getOrElse("") +
+ taskId.map("_TId_" + _).getOrElse("") +
+ taskAttemptNumber.map("_" + _).getOrElse("") +
+ upstreamCallerContext.map("_" + _).getOrElse(""))
+
+ private def prepareContext(context: String): String = {
+ lazy val len = SparkHadoopUtil.get.conf.getInt("hadoop.caller.context.max.size", 128)
+ if (context == null || context.length <= len) {
+ context
+ } else {
+ val finalContext = context.substring(0, len)
+ logWarning(s"Truncated Spark caller context from $context to $finalContext")
+ finalContext
+ }
+ }
+
+ def setCurrentContext(): Unit = {
+ if (callerContextSupported) {
+ try {
+ val callerContext = Utils.classForName("org.apache.hadoop.ipc.CallerContext")
+ val builder: Class[AnyRef] =
+ Utils.classForName("org.apache.hadoop.ipc.CallerContext$Builder")
+ val builderInst = builder.getConstructor(classOf[String]).newInstance(context)
+ val hdfsContext = builder.getMethod("build").invoke(builderInst)
+ callerContext.getMethod("setCurrent", callerContext).invoke(null, hdfsContext)
+ } catch {
+ case NonFatal(e) =>
+ logWarning("Fail to set Spark caller context", e)
+ }
+ }
+ }
+ }
+}
diff --git a/spark-extension/src/main/scala/org/apache/spark/sql/blaze/BlazeConverters.scala b/spark-extension/src/main/scala/org/apache/spark/sql/blaze/BlazeConverters.scala
index c9af4188..d0ebc60c 100644
--- a/spark-extension/src/main/scala/org/apache/spark/sql/blaze/BlazeConverters.scala
+++ b/spark-extension/src/main/scala/org/apache/spark/sql/blaze/BlazeConverters.scala
@@ -369,7 +369,7 @@ object BlazeConverters extends Logging {
case p =>
throw new NotImplementedError(
s"Cannot convert FileSourceScanExec tableIdentifier: ${tableIdentifier.getOrElse(
- "unknown")}, class: ${p.getClass.getName}")
+ "unknown")}, class: ${p.getClass.getName}")
}
}
@@ -734,8 +734,7 @@ object BlazeConverters extends Logging {
addRenameColumnsExec(convertToNative(exec.child))
case _ =>
if (needRenameColumns(exec.child)) {
- val newNames = exec.groupingExpressions.map(Util.getFieldNameByExprId) :+
- NativeAggBase.AGG_BUF_COLUMN_NAME
+ val newNames = exec.groupingExpressions.map(Util.getFieldNameByExprId)
Shims.get.createNativeRenameColumnsExec(convertToNative(exec.child), newNames)
} else {
convertToNative(exec.child)
@@ -791,8 +790,7 @@ object BlazeConverters extends Logging {
addRenameColumnsExec(convertToNative(exec.child))
case _ =>
if (needRenameColumns(exec.child)) {
- val newNames = exec.groupingExpressions.map(Util.getFieldNameByExprId) :+
- NativeAggBase.AGG_BUF_COLUMN_NAME
+ val newNames = exec.groupingExpressions.map(Util.getFieldNameByExprId)
Shims.get.createNativeRenameColumnsExec(convertToNative(exec.child), newNames)
} else {
convertToNative(exec.child)
@@ -845,8 +843,7 @@ object BlazeConverters extends Logging {
addRenameColumnsExec(convertToNative(child))
case _ =>
if (needRenameColumns(child)) {
- val newNames = exec.groupingExpressions.map(Util.getFieldNameByExprId) :+
- NativeAggBase.AGG_BUF_COLUMN_NAME
+ val newNames = exec.groupingExpressions.map(Util.getFieldNameByExprId)
Shims.get.createNativeRenameColumnsExec(convertToNative(child), newNames)
} else {
convertToNative(child)
diff --git a/spark-extension/src/main/scala/org/apache/spark/sql/blaze/BlazeSparkSessionExtension.scala b/spark-extension/src/main/scala/org/apache/spark/sql/blaze/BlazeSparkSessionExtension.scala
index 6db25c84..e26b9a7f 100644
--- a/spark-extension/src/main/scala/org/apache/spark/sql/blaze/BlazeSparkSessionExtension.scala
+++ b/spark-extension/src/main/scala/org/apache/spark/sql/blaze/BlazeSparkSessionExtension.scala
@@ -85,7 +85,7 @@ case class BlazeColumnarOverrides(sparkSession: SparkSession) extends ColumnarRu
dumpSimpleSparkPlanTreeNode(sparkPlanTransformed)
logInfo(s"Transformed spark plan after preColumnarTransitions:\n${sparkPlanTransformed
- .treeString(verbose = true, addSuffix = true)}")
+ .treeString(verbose = true, addSuffix = true)}")
// post-transform
Shims.get.postTransform(sparkPlanTransformed, sparkSession.sparkContext)
diff --git a/spark-extension/src/main/scala/org/apache/spark/sql/blaze/NativeHelper.scala b/spark-extension/src/main/scala/org/apache/spark/sql/blaze/NativeHelper.scala
index 42a3775d..4b8659f9 100644
--- a/spark-extension/src/main/scala/org/apache/spark/sql/blaze/NativeHelper.scala
+++ b/spark-extension/src/main/scala/org/apache/spark/sql/blaze/NativeHelper.scala
@@ -61,7 +61,8 @@ object NativeHelper extends Logging {
val nativeMemory: Long = {
val heapMemory = Runtime.getRuntime.maxMemory()
val offheapMemory = totalMemory - heapMemory
- logWarning(s"memory total: $totalMemory, onheap: $heapMemory, offheap: $offheapMemory")
+ logWarning(s"memory total: ${Utils.bytesToString(totalMemory)}, onheap: ${Utils.bytesToString(
+ heapMemory)}, offheap: ${Utils.bytesToString(offheapMemory)}")
offheapMemory
}
diff --git a/spark-extension/src/main/scala/org/apache/spark/sql/execution/blaze/plan/NativeAggBase.scala b/spark-extension/src/main/scala/org/apache/spark/sql/execution/blaze/plan/NativeAggBase.scala
index 13bbf877..2078117b 100644
--- a/spark-extension/src/main/scala/org/apache/spark/sql/execution/blaze/plan/NativeAggBase.scala
+++ b/spark-extension/src/main/scala/org/apache/spark/sql/execution/blaze/plan/NativeAggBase.scala
@@ -55,7 +55,7 @@ import org.apache.spark.sql.execution.blaze.plan.NativeAggBase.HashAgg
import org.apache.spark.sql.execution.blaze.plan.NativeAggBase.NativeAggrInfo
import org.apache.spark.sql.execution.blaze.plan.NativeAggBase.SortAgg
import org.apache.spark.sql.execution.metric.SQLMetrics
-import org.apache.spark.sql.types.BinaryType
+import org.apache.spark.sql.types.{BinaryType, BooleanType, DataType, LongType}
abstract class NativeAggBase(
execMode: AggExecMode,
@@ -143,9 +143,10 @@ abstract class NativeAggBase(
if (nativeAggrModes.contains(pb.AggMode.FINAL)) {
groupingExpressions.map(_.toAttribute) ++ aggregateAttributes
} else {
- groupingExpressions.map(_.toAttribute) :+
- AttributeReference(NativeAggBase.AGG_BUF_COLUMN_NAME, BinaryType, nullable = false)(
- ExprId.apply(NativeAggBase.AGG_BUF_COLUMN_EXPR_ID))
+ val aggBufferAttrs = nativeAggrInfos
+ .flatMap(_.aggBufferDataTypes)
+ .map(AttributeReference("", _, nullable = false)(ExprId.apply(0)))
+ groupingExpressions.map(_.toAttribute) ++ aggBufferAttrs
}
override def outputPartitioning: Partitioning =
@@ -205,38 +206,18 @@ abstract class NativeAggBase(
}
object NativeAggBase extends Logging {
-
- val AGG_BUF_COLUMN_EXPR_ID = 9223372036854775807L
- val AGG_BUF_COLUMN_NAME = s"#$AGG_BUF_COLUMN_EXPR_ID"
-
trait AggExecMode;
case object HashAgg extends AggExecMode
case object SortAgg extends AggExecMode
- case class NativeAggrPartialState(stateAttr: Attribute, arrowType: pb.ArrowType)
-
- object NativeAggrPartialState {
- def apply(
- aggrAttr: Attribute,
- stateFieldName: String,
- dataType: DataType,
- nullable: Boolean,
- arrowType: pb.ArrowType = null): NativeAggrPartialState = {
-
- val fieldName = s"${Util.getFieldNameByExprId(aggrAttr)}[$stateFieldName]"
- val stateAttr = AttributeReference(fieldName, dataType, nullable)(aggrAttr.exprId)
- NativeAggrPartialState(
- stateAttr,
- arrowType = Option(arrowType).getOrElse(NativeConverters.convertDataType(dataType)))
- }
- }
-
case class NativeAggrInfo(
mode: AggregateMode,
nativeAggrs: Seq[pb.PhysicalExprNode],
+ aggBufferDataTypes: Seq[DataType],
outputAttr: Attribute)
def getNativeAggrInfo(aggr: AggregateExpression, aggrAttr: Attribute): NativeAggrInfo = {
+ val aggBufferDataTypes = computeNativeAggBufferDataTypes(aggr.aggregateFunction)
val reducedAggr = AggregateExpression(
aggr.aggregateFunction
.mapChildren(e => createPlaceholder(e))
@@ -249,12 +230,17 @@ object NativeAggBase extends Logging {
aggr.mode match {
case Partial =>
- NativeAggrInfo(aggr.mode, NativeConverters.convertAggregateExpr(aggr) :: Nil, outputAttr)
+ NativeAggrInfo(
+ aggr.mode,
+ NativeConverters.convertAggregateExpr(aggr) :: Nil,
+ aggBufferDataTypes,
+ outputAttr)
case PartialMerge | Final =>
NativeAggrInfo(
aggr.mode,
NativeConverters.convertAggregateExpr(reducedAggr) :: Nil,
+ aggBufferDataTypes,
outputAttr)
case Complete =>
@@ -306,4 +292,17 @@ object NativeAggBase extends Logging {
}
findRecursive(exec.children.head)
}
+
+ def computeNativeAggBufferDataTypes(aggr: AggregateFunction): Seq[DataType] = {
+ aggr match {
+ case _: Count => Seq(LongType)
+ case f: Max => Seq(f.dataType)
+ case f: Min => Seq(f.dataType)
+ case f: Sum => Seq(f.dataType)
+ case f: Average => Seq(f.dataType, LongType)
+ case f @ First(_, true) => Seq(f.dataType)
+ case f @ First(_, false) => Seq(f.dataType, BooleanType)
+ case _ => Seq(BinaryType)
+ }
+ }
}
(auron) 15/16: supports spark.blaze.enableScanFallback
Posted by ri...@apache.org.
This is an automated email from the ASF dual-hosted git repository.
richox pushed a commit to branch dev-v6.0.0-parallel-scan-kdev-build
in repository https://gitbox.apache.org/repos/asf/auron.git
commit ea2dd42b48f3ae30ea7d2277ecb15227a428426c
Author: zhangli20 <ri...@gmail.com>
AuthorDate: Sat Mar 7 16:15:57 2026 +0800
supports spark.blaze.enableScanFallback
---
.idea/vcs.xml | 3 ---
.../src/main/scala/org/apache/spark/sql/blaze/ShimsImpl.scala | 7 +++++--
.../src/main/scala/org/apache/spark/sql/blaze/ShimsImpl.scala | 7 +++++--
.../org/apache/spark/sql/execution/blaze/plan/NativeAggExec.scala | 1 -
.../src/main/java/org/apache/spark/sql/blaze/BlazeConf.java | 2 ++
.../src/main/scala/org/apache/spark/sql/blaze/NativeHelper.scala | 3 +--
6 files changed, 13 insertions(+), 10 deletions(-)
diff --git a/.idea/vcs.xml b/.idea/vcs.xml
index bd37bc65..02158c9e 100644
--- a/.idea/vcs.xml
+++ b/.idea/vcs.xml
@@ -12,8 +12,5 @@
</component>
<component name="VcsDirectoryMappings">
<mapping directory="$PROJECT_DIR$" vcs="Git" />
- <mapping directory="$PROJECT_DIR$/thirdparty/arrow-datafusion" vcs="Git" />
- <mapping directory="$PROJECT_DIR$/thirdparty/arrow-rs" vcs="Git" />
- <mapping directory="$PROJECT_DIR$/thirdparty/datafusion-orc" vcs="Git" />
</component>
</project>
\ No newline at end of file
diff --git a/spark-extension-shims-spark241kwaiae/src/main/scala/org/apache/spark/sql/blaze/ShimsImpl.scala b/spark-extension-shims-spark241kwaiae/src/main/scala/org/apache/spark/sql/blaze/ShimsImpl.scala
index 60b203e6..26820773 100644
--- a/spark-extension-shims-spark241kwaiae/src/main/scala/org/apache/spark/sql/blaze/ShimsImpl.scala
+++ b/spark-extension-shims-spark241kwaiae/src/main/scala/org/apache/spark/sql/blaze/ShimsImpl.scala
@@ -746,7 +746,7 @@ case class BlazeRuleEngine(sparkSession: SparkSession) extends Rule[LogicalPlan]
plan.foreachUp {
case p @ LogicalRelation(fsRelation: HadoopFsRelation, _, _, _) =>
// non parquet table rule
- if (!(
+ if (!BlazeConf.ENABLE_SCAN_FALLBACK.booleanConf() && !(
BlazeConverters.enableScanParquet && fsRelation.fileFormat.isInstanceOf[ParquetFileFormat]
|| BlazeConverters.enableScanOrc && fsRelation.fileFormat.isInstanceOf[OrcFileFormat]
)) {
@@ -775,7 +775,9 @@ case class BlazeRuleEngine(sparkSession: SparkSession) extends Rule[LogicalPlan]
}
case h: HiveTableRelation =>
- turnOffBlazeWithReason(h.conf, BlazeMissPatterns.NonParquetFormat)
+ if (!BlazeConf.ENABLE_SCAN_FALLBACK.booleanConf()) {
+ turnOffBlazeWithReason(h.conf, BlazeMissPatterns.NonParquetFormat)
+ }
case _ =>
}
@@ -783,6 +785,7 @@ case class BlazeRuleEngine(sparkSession: SparkSession) extends Rule[LogicalPlan]
}
private def turnOffBlazeWithReason(planConf: SQLConf, blazeMissPattern: String): Unit = {
+ logWarning(s"Turn off Blaze due to: $blazeMissPattern")
planConf.setConf(blazeEnabledKey, false)
sparkSession.sparkContext.conf
.set(BlazeRuleEngine.blazeMissPatterns, blazeMissPattern)
diff --git a/spark-extension-shims-spark3/src/main/scala/org/apache/spark/sql/blaze/ShimsImpl.scala b/spark-extension-shims-spark3/src/main/scala/org/apache/spark/sql/blaze/ShimsImpl.scala
index 8b74e735..2b6ca98d 100644
--- a/spark-extension-shims-spark3/src/main/scala/org/apache/spark/sql/blaze/ShimsImpl.scala
+++ b/spark-extension-shims-spark3/src/main/scala/org/apache/spark/sql/blaze/ShimsImpl.scala
@@ -1063,7 +1063,7 @@ case class BlazeRuleEngine(sparkSession: SparkSession) extends Rule[LogicalPlan]
plan.foreachUp {
case p @ LogicalRelation(fsRelation: HadoopFsRelation, _, _, _) =>
// non parquet table rule
- if (!(
+ if (!BlazeConf.ENABLE_SCAN_FALLBACK.booleanConf() && !(
BlazeConverters.enableScanParquet && fsRelation.fileFormat.isInstanceOf[ParquetFileFormat]
|| BlazeConverters.enableScanOrc && fsRelation.fileFormat.isInstanceOf[OrcFileFormat]
)) {
@@ -1092,7 +1092,9 @@ case class BlazeRuleEngine(sparkSession: SparkSession) extends Rule[LogicalPlan]
}
case h: HiveTableRelation =>
- turnOffBlazeWithReason(h.conf, BlazeMissPatterns.NonParquetFormat)
+ if (!BlazeConf.ENABLE_SCAN_FALLBACK.booleanConf()) {
+ turnOffBlazeWithReason(h.conf, BlazeMissPatterns.NonParquetFormat)
+ }
case _ =>
}
@@ -1100,6 +1102,7 @@ case class BlazeRuleEngine(sparkSession: SparkSession) extends Rule[LogicalPlan]
}
private def turnOffBlazeWithReason(planConf: SQLConf, blazeMissPattern: String): Unit = {
+ logWarning(s"Turn off Blaze due to: $blazeMissPattern")
planConf.setConf(blazeEnabledKey, false)
sparkSession.sparkContext.conf
.set(BlazeRuleEngine.blazeMissPatterns, blazeMissPattern)
diff --git a/spark-extension-shims-spark3/src/main/scala/org/apache/spark/sql/execution/blaze/plan/NativeAggExec.scala b/spark-extension-shims-spark3/src/main/scala/org/apache/spark/sql/execution/blaze/plan/NativeAggExec.scala
index 6cfbbc36..1cee51b4 100644
--- a/spark-extension-shims-spark3/src/main/scala/org/apache/spark/sql/execution/blaze/plan/NativeAggExec.scala
+++ b/spark-extension-shims-spark3/src/main/scala/org/apache/spark/sql/execution/blaze/plan/NativeAggExec.scala
@@ -22,7 +22,6 @@ import org.apache.spark.sql.catalyst.expressions.AttributeReference
import org.apache.spark.sql.catalyst.expressions.Expression
import org.apache.spark.sql.catalyst.expressions.NamedExpression
import org.apache.spark.sql.execution.aggregate.BaseAggregateExec
-import org.apache.spark.sql.execution.auron.plan.NativeAggBase.AggExecMode
import org.apache.spark.sql.execution.blaze.plan.NativeAggBase.AggExecMode
import org.apache.spark.sql.execution.SparkPlan
import org.apache.spark.sql.types.BinaryType
diff --git a/spark-extension/src/main/java/org/apache/spark/sql/blaze/BlazeConf.java b/spark-extension/src/main/java/org/apache/spark/sql/blaze/BlazeConf.java
index 5136b8b5..76f90653 100644
--- a/spark-extension/src/main/java/org/apache/spark/sql/blaze/BlazeConf.java
+++ b/spark-extension/src/main/java/org/apache/spark/sql/blaze/BlazeConf.java
@@ -115,6 +115,8 @@ public enum BlazeConf {
// number of parallel scan files
NUM_PARALLEL_SCAN_FILES("spark.blaze.numParallelScanFiles", 4),
+ ENABLE_SCAN_FALLBACK("spark.blaze.enableScanFallback", false),
+
NATIVE_LOG_LEVEL("spark.blaze.native.log.level", "info");
public final String key;
diff --git a/spark-extension/src/main/scala/org/apache/spark/sql/blaze/NativeHelper.scala b/spark-extension/src/main/scala/org/apache/spark/sql/blaze/NativeHelper.scala
index 4b8659f9..5ea13b36 100644
--- a/spark-extension/src/main/scala/org/apache/spark/sql/blaze/NativeHelper.scala
+++ b/spark-extension/src/main/scala/org/apache/spark/sql/blaze/NativeHelper.scala
@@ -16,7 +16,6 @@
package org.apache.spark.sql.blaze
import scala.collection.immutable.TreeMap
-
import org.apache.hadoop.security.UserGroupInformation
import org.apache.spark.Partition
import org.apache.spark.SparkConf
@@ -24,13 +23,13 @@ import org.apache.spark.SparkContext
import org.apache.spark.SparkEnv
import org.apache.spark.TaskContext
import org.blaze.protobuf.PhysicalPlanNode
-
import org.apache.spark.internal.Logging
import org.apache.spark.internal.config
import org.apache.spark.sql.catalyst.InternalRow
import org.apache.spark.sql.execution.SparkPlan
import org.apache.spark.sql.execution.metric.SQLMetric
import org.apache.spark.sql.execution.metric.SQLMetrics
+import org.apache.spark.util.Utils
object NativeHelper extends Logging {
val currentUser: UserGroupInformation = UserGroupInformation.getCurrentUser
(auron) 06/16: adapt to celeborn-0.6.0 with some special patchs
Posted by ri...@apache.org.
This is an automated email from the ASF dual-hosted git repository.
richox pushed a commit to branch dev-v6.0.0-parallel-scan-kdev-build
in repository https://gitbox.apache.org/repos/asf/auron.git
commit d446a83983ef383b1622554ecdba933e824d2b65
Author: zhuangxian <zh...@kuaishou.com>
AuthorDate: Tue Dec 9 02:31:20 2025 +0000
adapt to celeborn-0.6.0 with some special patchs
KDev_MR_link：https://ksurl.cn/cMld7QJI
---
pom.xml | 4 ++--
.../shuffle/celeborn/CelebornPartitionWriter.scala | 22 +++++++++++++++++++++-
2 files changed, 23 insertions(+), 3 deletions(-)
diff --git a/pom.xml b/pom.xml
index b455cdac..57140bf3 100644
--- a/pom.xml
+++ b/pom.xml
@@ -21,7 +21,7 @@
<arrowVersion>16.0.0</arrowVersion>
<protobufVersion>3.21.9</protobufVersion>
<paimonVersion>1.1.1</paimonVersion>
- <celebornVersion>0.6.0</celebornVersion>
+ <celebornVersion>0.6.2</celebornVersion>
</properties>
<dependencyManagement>
@@ -416,7 +416,7 @@
<id>spark-3.5</id>
<properties>
<shimName>spark-3.5</shimName>
- <pkgSuffix>-kwai</pkgSuffix>
+ <pkgSuffix>-kwai-adapt-cele060</pkgSuffix>
<shimPkg>spark-extension-shims-spark3</shimPkg>
<javaVersion>1.8</javaVersion>
<scalaVersion>2.12</scalaVersion>
diff --git a/spark-extension-shims-spark3/src/main/scala/org/apache/spark/sql/execution/blaze/shuffle/celeborn/CelebornPartitionWriter.scala b/spark-extension-shims-spark3/src/main/scala/org/apache/spark/sql/execution/blaze/shuffle/celeborn/CelebornPartitionWriter.scala
index 88e2eb6f..a60ec70e 100644
--- a/spark-extension-shims-spark3/src/main/scala/org/apache/spark/sql/execution/blaze/shuffle/celeborn/CelebornPartitionWriter.scala
+++ b/spark-extension-shims-spark3/src/main/scala/org/apache/spark/sql/execution/blaze/shuffle/celeborn/CelebornPartitionWriter.scala
@@ -64,7 +64,27 @@ class CelebornPartitionWriter(
override def close(success: Boolean): Unit = {
val waitStartTime = System.nanoTime()
if (success) {
- shuffleClient.mapperEnd(shuffleId, mapId, encodedAttemptId, numMappers)
+ try {
+ // Try new API with numPartitions parameter (Celeborn 0.6+)
+ shuffleClient.mapperEnd(shuffleId, mapId, encodedAttemptId, numMappers, numPartitions)
+ } catch {
+ case _: NoSuchMethodError =>
+ // Fallback to old API without numPartitions parameter (Celeborn 0.5.x)
+ // Use reflection to avoid compilation error with new version
+ try {
+ val method = shuffleClient.getClass.getMethod(
+ "mapperEnd",
+ classOf[Int], classOf[Int], classOf[Int], classOf[Int])
+ method.invoke(shuffleClient,
+ Int.box(shuffleId),
+ Int.box(mapId),
+ Int.box(encodedAttemptId),
+ Int.box(numMappers))
+ } catch {
+ case e: Exception =>
+ logWarning(s"mapperEnd failed with error: ${e.getMessage}", e)
+ }
+ }
}
shuffleClient.cleanup(shuffleId, mapId, encodedAttemptId)
metrics.incWriteTime(System.nanoTime() - waitStartTime)
(auron) 01/16: Shade more packages (#1121)
Posted by ri...@apache.org.
This is an automated email from the ASF dual-hosted git repository.
richox pushed a commit to branch dev-v6.0.0-parallel-scan-kdev-build
in repository https://gitbox.apache.org/repos/asf/auron.git
commit 494645c0cf1907a7dd0cf63cd06767333af77b35
Author: Fei Wang <fw...@ebay.com>
AuthorDate: Thu Aug 7 02:49:23 2025 -0700
Shade more packages (#1121)
* Shade
* typo
* comments
---
dev/mvn-build-helper/assembly/pom.xml | 36 ++++++++++++++++++++++++++++-------
1 file changed, 29 insertions(+), 7 deletions(-)
diff --git a/dev/mvn-build-helper/assembly/pom.xml b/dev/mvn-build-helper/assembly/pom.xml
index d522cfc4..dae047e9 100644
--- a/dev/mvn-build-helper/assembly/pom.xml
+++ b/dev/mvn-build-helper/assembly/pom.xml
@@ -12,6 +12,10 @@
<artifactId>assembly</artifactId>
<packaging>jar</packaging>
+ <properties>
+ <blaze.shade.packageName>blaze</blaze.shade.packageName>
+ </properties>
+
<dependencies>
<dependency>
<groupId>org.blaze</groupId>
@@ -61,29 +65,47 @@
<exclude>org/apache/commons/codec/**/*</exclude>
<exclude>org/apache/commons/compress/**/*</exclude>
<exclude>org/slf4j/**/*</exclude>
- <excludes>META-INF/*.SF</excludes>
- <excludes>META-INF/*.DSA</excludes>
- <excludes>META-INF/*.RSA</excludes>
- <excludes>**/*.semanticdb</excludes>
+ <exclude>META-INF/*.SF</exclude>
+ <exclude>META-INF/*.DSA</exclude>
+ <exclude>META-INF/*.RSA</exclude>
+ <exclude>META-INF/**/module-info.class</exclude>
+ <exclude>**/*.semanticdb</exclude>
</excludes>
</filter>
</filters>
<relocations>
<relocation>
<pattern>com.google.flatbuffers</pattern>
- <shadedPattern>blaze.com.google.flatbuffers</shadedPattern>
+ <shadedPattern>${blaze.shade.packageName}.com.google.flatbuffers</shadedPattern>
</relocation>
<relocation>
<pattern>com.google.protobuf</pattern>
- <shadedPattern>blaze.com.google.protobuf</shadedPattern>
+ <shadedPattern>${blaze.shade.packageName}.com.google.protobuf</shadedPattern>
</relocation>
<relocation>
<pattern>org.apache.arrow</pattern>
- <shadedPattern>blaze.org.apache.arrow</shadedPattern>
+ <shadedPattern>${blaze.shade.packageName}.org.apache.arrow</shadedPattern>
+ <!--arrow's C and dataset wrapper refers to the original class path, so we should not relocate here-->
<excludes>
<exclude>org.apache.arrow.c.**</exclude>
</excludes>
</relocation>
+ <relocation>
+ <pattern>javax.annotation</pattern>
+ <shadedPattern>${blaze.shade.packageName}.javax.annotation</shadedPattern>
+ </relocation>
+ <relocation>
+ <pattern>net.bytebuddy</pattern>
+ <shadedPattern>${blaze.shade.packageName}.net.bytebuddy</shadedPattern>
+ </relocation>
+ <relocation>
+ <pattern>org.checkerframework</pattern>
+ <shadedPattern>${blaze.shade.packageName}.org.checkerframework</shadedPattern>
+ </relocation>
+ <relocation>
+ <pattern>scala.compat</pattern>
+ <shadedPattern>${blaze.shade.packageName}.scala.compat</shadedPattern>
+ </relocation>
</relocations>
<archive>
<addMavenDescriptor>false</addMavenDescriptor>
(auron) 09/16: introduce BlazeRuleEngine from spark241kwaiae to spark3
Posted by ri...@apache.org.
This is an automated email from the ASF dual-hosted git repository.
richox pushed a commit to branch dev-v6.0.0-parallel-scan-kdev-build
in repository https://gitbox.apache.org/repos/asf/auron.git
commit c44088557602444e766a05c918780f7ff3b43dfe
Author: zhangli20 <ri...@gmail.com>
AuthorDate: Thu Dec 11 20:41:28 2025 +0800
introduce BlazeRuleEngine from spark241kwaiae to spark3
---
pom.xml | 2 +-
.../org/apache/spark/sql/blaze/ShimsImpl.scala | 81 ++++++++++++++++++++--
2 files changed, 78 insertions(+), 5 deletions(-)
diff --git a/pom.xml b/pom.xml
index 0301c25f..1a186ce6 100644
--- a/pom.xml
+++ b/pom.xml
@@ -416,7 +416,7 @@
<id>spark-3.5</id>
<properties>
<shimName>spark-3.5</shimName>
- <pkgSuffix>-kwai-adapt-cele060-fix-non-partition-len</pkgSuffix>
+ <pkgSuffix>-kwai-hotfix-rule-engine</pkgSuffix>
<shimPkg>spark-extension-shims-spark3</shimPkg>
<javaVersion>1.8</javaVersion>
<scalaVersion>2.12</scalaVersion>
diff --git a/spark-extension-shims-spark3/src/main/scala/org/apache/spark/sql/blaze/ShimsImpl.scala b/spark-extension-shims-spark3/src/main/scala/org/apache/spark/sql/blaze/ShimsImpl.scala
index 44ce358e..8b74e735 100644
--- a/spark-extension-shims-spark3/src/main/scala/org/apache/spark/sql/blaze/ShimsImpl.scala
+++ b/spark-extension-shims-spark3/src/main/scala/org/apache/spark/sql/blaze/ShimsImpl.scala
@@ -17,7 +17,6 @@ package org.apache.spark.sql.blaze
import java.io.File
import java.util.UUID
-
import org.apache.commons.lang3.reflect.FieldUtils
import org.apache.spark.OneToOneDependency
import org.apache.spark.ShuffleDependency
@@ -25,6 +24,7 @@ import org.apache.spark.SparkEnv
import org.apache.spark.SparkException
import org.apache.spark.TaskContext
import org.apache.spark.internal.Logging
+import org.apache.spark.internal.config.OptionalConfigEntry
import org.apache.spark.rdd.RDD
import org.apache.spark.scheduler.MapStatus
import org.apache.spark.shuffle.IndexShuffleBlockResolver
@@ -37,7 +37,7 @@ import org.apache.spark.sql.SparkSession
import org.apache.spark.sql.blaze.BlazeConverters.ForceNativeExecutionWrapperBase
import org.apache.spark.sql.blaze.NativeConverters.NativeExprWrapperBase
import org.apache.spark.sql.catalyst.InternalRow
-import org.apache.spark.sql.catalyst.catalog.CatalogTable
+import org.apache.spark.sql.catalyst.catalog.{CatalogTable, HiveTableRelation}
import org.apache.spark.sql.catalyst.expressions.Attribute
import org.apache.spark.sql.catalyst.expressions.Expression
import org.apache.spark.sql.catalyst.expressions.Generator
@@ -101,7 +101,7 @@ import org.apache.spark.sql.execution.blaze.plan._
import org.apache.spark.sql.execution.blaze.shuffle.RssPartitionWriterBase
import org.apache.spark.sql.execution.blaze.shuffle.celeborn.BlazeCelebornShuffleManager
import org.apache.spark.sql.execution.blaze.shuffle.BlazeBlockStoreShuffleReaderBase
-import org.apache.spark.sql.execution.datasources.PartitionedFile
+import org.apache.spark.sql.execution.datasources.{HadoopFsRelation, LogicalRelation, PartitionedFile}
import org.apache.spark.sql.execution.exchange.{BroadcastExchangeLike, ReusedExchangeExec}
import org.apache.spark.sql.execution.joins.blaze.plan.NativeBroadcastJoinExec
import org.apache.spark.sql.execution.joins.blaze.plan.NativeShuffledHashJoinExecProvider
@@ -112,12 +112,17 @@ import org.apache.spark.sql.types.DataType
import org.apache.spark.sql.types.IntegerType
import org.apache.spark.sql.types.StringType
import org.apache.spark.sql.SparkSessionExtensions
+import org.apache.spark.sql.catalyst.plans.logical.LogicalPlan
+import org.apache.spark.sql.catalyst.rules.Rule
import org.apache.spark.storage.BlockManagerId
import org.apache.spark.storage.FileSegment
import org.apache.spark.sql.execution.blaze.shuffle.uniffle.BlazeUniffleShuffleManager
+import org.apache.spark.sql.execution.datasources.orc.OrcFileFormat
+import org.apache.spark.sql.execution.datasources.parquet.ParquetFileFormat
+import org.apache.spark.sql.internal.SQLConf
import org.apache.spark.storage.ShuffleDataBlockId
import org.apache.spark.util.ExternalBlockStoreUtils
-import org.blaze.{protobuf => pb, sparkver}
+import org.blaze.{sparkver, protobuf => pb}
class ShimsImpl extends Shims with Logging {
@@ -156,6 +161,10 @@ class ShimsImpl extends Shims with Logging {
logWarning(s"${BlazeConf.FORCE_SHUFFLED_HASH_JOIN.key} is not supported in $shimVersion")
}
}
+
+ extension.injectOptimizerRule(sparkSession => {
+ BlazeRuleEngine(sparkSession)
+ })
}
override def createConvertToNativeExec(child: SparkPlan): ConvertToNativeBase =
@@ -1046,3 +1055,67 @@ case class NativeExprWrapper(
@sparkver("3.2 / 3.3 / 3.4 / 3.5")
override def withNewChildrenInternal(newChildren: IndexedSeq[Expression]): Expression = copy()
}
+
+case class BlazeRuleEngine(sparkSession: SparkSession) extends Rule[LogicalPlan] with Logging {
+ import BlazeSparkSessionExtension._
+
+ override def apply(plan: LogicalPlan): LogicalPlan = {
+ plan.foreachUp {
+ case p @ LogicalRelation(fsRelation: HadoopFsRelation, _, _, _) =>
+ // non parquet table rule
+ if (!(
+ BlazeConverters.enableScanParquet && fsRelation.fileFormat.isInstanceOf[ParquetFileFormat]
+ || BlazeConverters.enableScanOrc && fsRelation.fileFormat.isInstanceOf[OrcFileFormat]
+ )) {
+ turnOffBlazeWithReason(p.conf, BlazeMissPatterns.NonParquetFormat)
+ }
+
+ // read encrypted table rule
+ val readEncryptedTableEnable = sparkSession.sparkContext.conf
+ .getBoolean("spark.hive.exist.read.encrypted.table", defaultValue = false)
+ if (readEncryptedTableEnable) {
+ turnOffBlazeWithReason(p.conf, BlazeMissPatterns.ReadEncryptedTable)
+ }
+
+ // skip scan dp_dd.*** table because parquet statics don't get min/max info
+ if (p.catalogTable.map(_.identifier.unquotedString).getOrElse("").contains("dp_dd")) {
+ turnOffBlazeWithReason(p.conf, BlazeMissPatterns.ReadHbaseTable)
+ }
+
+ // skip scan offline_attribution_mover_v2
+ // issue: https://team.corp.kuaishou.com/task/T6538018
+ if (p.catalogTable
+ .map(_.identifier.unquotedString)
+ .getOrElse("")
+ .contains("offline_attribution_mover_v2")) {
+ turnOffBlazeWithReason(p.conf, BlazeMissPatterns.ReadBlacklistedTable)
+ }
+
+ case h: HiveTableRelation =>
+ turnOffBlazeWithReason(h.conf, BlazeMissPatterns.NonParquetFormat)
+
+ case _ =>
+ }
+ plan
+ }
+
+ private def turnOffBlazeWithReason(planConf: SQLConf, blazeMissPattern: String): Unit = {
+ planConf.setConf(blazeEnabledKey, false)
+ sparkSession.sparkContext.conf
+ .set(BlazeRuleEngine.blazeMissPatterns, blazeMissPattern)
+ }
+
+ object BlazeMissPatterns extends Enumeration {
+ val NonParquetFormat = "NonParquetFormat"
+ val ReadEncryptedTable = "ReadEncryptedTable"
+ val ReadHbaseTable = "ReadHbaseTable"
+ val ReadBlacklistedTable = "ReadBlacklistedTable"
+ }
+}
+
+object BlazeRuleEngine {
+ lazy val blazeMissPatterns: OptionalConfigEntry[String] = SQLConf
+ .buildConf("spark.blaze.blazeMissPatterns")
+ .stringConf
+ .createOptional
+}
(auron) 16/16: add kdev-build-deploy.sh
Posted by ri...@apache.org.
This is an automated email from the ASF dual-hosted git repository.
richox pushed a commit to branch dev-v6.0.0-parallel-scan-kdev-build
in repository https://gitbox.apache.org/repos/asf/auron.git
commit 15b9fcc2d68df05d39a9ebed0f70af39150df052
Author: zhangli20 <ri...@gmail.com>
AuthorDate: Thu Mar 19 11:12:32 2026 +0800
add kdev-build-deploy.sh
---
kdev-build-deploy.sh | 3 +++
.../org/apache/spark/sql/execution/blaze/plan/NativeAggBase.scala | 2 +-
2 files changed, 4 insertions(+), 1 deletion(-)
diff --git a/kdev-build-deploy.sh b/kdev-build-deploy.sh
new file mode 100755
index 00000000..46e0a5b8
--- /dev/null
+++ b/kdev-build-deploy.sh
@@ -0,0 +1,3 @@
+#!/bin/bash
+
+mvn package -Pspark-3.5 -Prelease -Dmaven.test.skip -DskipTests
diff --git a/spark-extension/src/main/scala/org/apache/spark/sql/execution/blaze/plan/NativeAggBase.scala b/spark-extension/src/main/scala/org/apache/spark/sql/execution/blaze/plan/NativeAggBase.scala
index 5824fac3..26bcd9d2 100644
--- a/spark-extension/src/main/scala/org/apache/spark/sql/execution/blaze/plan/NativeAggBase.scala
+++ b/spark-extension/src/main/scala/org/apache/spark/sql/execution/blaze/plan/NativeAggBase.scala
@@ -227,7 +227,7 @@ object NativeAggBase extends Logging {
aggr.mode,
aggr.isDistinct)
val outputAttr =
- AttributeReference(Util.getFieldNameByExprId(aggrAttr), aggrAttr.dataType, aggr.nullable)(
+ AttributeReference(Util.getFieldNameByExprId(aggrAttr), aggrAttr.dataType, aggrAttr.nullable)(
aggrAttr.exprId)
aggr.mode match {