refactor(executor): add max rows/bytes limit when schedule executor task

Author: dqhl76

src/common/base/src/runtime/datablock_limit.rs

@@ -0,0 +1,98 @@
+// Copyright 2021 Datafuse Labs
+//
+// 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.
+
+/// An indicator to limit the size of datablocks transferred between processor
+#[derive(Debug, Copy, Clone)]
+pub struct DataBlockLimit {
+    /// The maximum number of rows allowed in a data block
+    pub max_rows: usize,
+    /// The maximum size in bytes allowed for a data block
+    pub max_bytes: usize,
+}
+
+impl DataBlockLimit {
+    pub fn new(mut max_rows: usize, mut max_bytes: usize) -> Self {
+        // If max_rows or max_bytes is set to 0, we treat it as no limit
+        if max_rows == 0 {
+            max_rows = usize::MAX;
+        }
+        if max_bytes == 0 {
+            max_bytes = usize::MAX;
+        }
+
+        Self {
+            max_rows,
+            max_bytes,
+        }
+    }
+
+    /// Calculate the number of rows to take based on both row and byte limits
+    pub fn rows_to_take(&self, total_rows: usize, total_bytes: usize) -> usize {
+        debug_assert!(total_rows > 0);
+        // check row limit
+        let rows_by_limit = self.max_rows.min(total_rows);
+
+        // check byte limit
+        let average_bytes_per_row = total_bytes / total_rows;
+        let rows_by_bytes = if average_bytes_per_row > 0 {
+            // TODO: this 1 may be not suitable
+            (self.max_bytes / average_bytes_per_row).max(1)
+        } else {
+            // Avoid division by zero
+            total_rows
+        };
+
+        // the minimum of the two limits
+        rows_by_limit.min(rows_by_bytes)
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn test_rows_to_take_basic() {
+        let limit = DataBlockLimit::new(1000, 1_000_000);
+
+        // Case 1: Both limits not exceeded
+        let rows = limit.rows_to_take(500, 500_000);
+        assert_eq!(rows, 500);
+
+        // Case 2: Row limit exceeded
+        let rows = limit.rows_to_take(2000, 500_000);
+        assert_eq!(rows, 1000);
+
+        // Case 3: Byte limit exceeded
+        // 2_000_000 bytes / 1000 rows = 2000 bytes per row
+        // 1_000_000 bytes / 2000 bytes per row = 500 rows
+        let rows = limit.rows_to_take(1000, 2_000_000);
+        assert_eq!(rows, 500);
+
+        let limit = DataBlockLimit::new(1000, 1_000_000);
+
+        // Both limits exceeded, byte limit is more restrictive
+        // 4_000_000 bytes / 2000 rows = 2000 bytes per row
+        // 1_000_000 bytes / 2000 bytes per row = 500 rows
+        let rows = limit.rows_to_take(2000, 4_000_000);
+        assert_eq!(rows, 500);
+
+        // Both limits exceeded, row limit is more restrictive
+        // 1_500_000 bytes / 2000 rows = 750 bytes per row
+        // 1_000_000 bytes / 750 bytes per row = 1333 rows
+        // But row limit is 1000, so we take 1000
+        let rows = limit.rows_to_take(2000, 1_500_000);
+        assert_eq!(rows, 1000);
+    }
+}

src/common/base/src/runtime/mod.rs

@@ -14,6 +14,7 @@
 
 mod backtrace;
 mod catch_unwind;
+mod datablock_limit;
 mod defer;
 pub mod error_info;
 mod executor_stats;
@@ -35,6 +36,7 @@ pub use backtrace::AsyncTaskItem;
 pub use catch_unwind::catch_unwind;
 pub use catch_unwind::drop_guard;
 pub use catch_unwind::CatchUnwindFuture;
+pub use datablock_limit::DataBlockLimit;
 pub use defer::defer;
 pub use executor_stats::ExecutorStats;
 pub use executor_stats::ExecutorStatsSlot;

src/common/base/src/runtime/runtime_tracker.rs

@@ -54,6 +54,7 @@ use concurrent_queue::ConcurrentQueue;
 use log::LevelFilter;
 use pin_project_lite::pin_project;
 
+use crate::runtime::datablock_limit::DataBlockLimit;
 use crate::runtime::memory::GlobalStatBuffer;
 use crate::runtime::memory::MemStat;
 use crate::runtime::metrics::ScopedRegistry;
@@ -144,6 +145,7 @@ pub struct TrackingPayload {
     pub workload_group_resource: Option<Arc<WorkloadGroupResource>>,
     pub perf_enabled: bool,
     pub process_rows: AtomicUsize,
+    pub datablock_limit: Option<Arc<DataBlockLimit>>,
 }
 
 impl Clone for TrackingPayload {
@@ -161,6 +163,7 @@ impl Clone for TrackingPayload {
             process_rows: AtomicUsize::new(
                 self.process_rows.load(std::sync::atomic::Ordering::SeqCst),
             ),
+            datablock_limit: self.datablock_limit.clone(),
         }
     }
 }
@@ -243,6 +246,7 @@ impl ThreadTracker {
                 workload_group_resource: None,
                 perf_enabled: false,
                 process_rows: AtomicUsize::new(0),
+                datablock_limit: None,
             }),
         }
     }
@@ -369,6 +373,20 @@ impl ThreadTracker {
             })
             .unwrap_or(0)
     }
+
+    pub fn datablock_limit() -> Option<&'static Arc<DataBlockLimit>> {
+        TRACKER
+            .try_with(|tracker| {
+                tracker
+                    .borrow()
+                    .payload
+                    .datablock_limit
+                    .as_ref()
+                    .map(|v| unsafe { std::mem::transmute(v) })
+            })
+            .ok()
+            .and_then(|x| x)
+    }
 }
 
 pin_project! {

src/query/expression/src/block.rs

@@ -547,6 +547,18 @@ impl DataBlock {
         }
     }
 
+    pub fn split_at(&self, mid: usize) -> (Self, Self) {
+        assert!(
+            mid <= self.num_rows,
+            "split point {} out of len {}",
+            mid,
+            self.num_rows
+        );
+        let first = self.slice(0..mid);
+        let second = self.slice(mid..self.num_rows);
+        (first, second)
+    }
+
     pub fn split_by_rows(&self, max_rows_per_block: usize) -> (Vec<Self>, Option<Self>) {
         let mut res = Vec::with_capacity(self.num_rows / max_rows_per_block);
         let mut offset = 0;

src/query/pipeline/core/src/processors/port.rs

@@ -19,8 +19,9 @@ use std::sync::Arc;
 use databend_common_base::runtime::drop_guard;
 use databend_common_base::runtime::profile::Profile;
 use databend_common_base::runtime::profile::ProfileStatisticsName;
-use databend_common_base::runtime::ExecutorStats;
+use databend_common_base::runtime::DataBlockLimit;
 use databend_common_base::runtime::QueryTimeSeriesProfile;
+use databend_common_base::runtime::ThreadTracker;
 use databend_common_base::runtime::TimeSeriesProfileName;
 use databend_common_exception::Result;
 use databend_common_expression::DataBlock;
@@ -186,17 +187,57 @@ impl InputPort {
 
     #[inline(always)]
     pub fn pull_data(&self) -> Option<Result<DataBlock>> {
+        let datablock_limit = ThreadTracker::datablock_limit();
+        if let Some(limit) = datablock_limit {
+            self.pull_partial_data(limit)
+        } else {
+            self.pull_all_data()
+        }
+    }
+
+    /// Pull entire available datablock without any limit.
+    ///
+    /// This method provides the original behavior with minimal overhead
+    /// when no datablock limit is configured.
+    #[inline(always)]
+    pub fn pull_all_data(&self) -> Option<Result<DataBlock>> {
         unsafe {
             UpdateTrigger::update_input(&self.update_trigger);
             let unset_flags = HAS_DATA | NEED_DATA;
             match self.shared.swap(std::ptr::null_mut(), 0, unset_flags) {
+                address if address.is_null() => None,
+                address => Some((*Box::from_raw(address)).0),
+            }
+        }
+    }
+
+    /// Pull datablock with a size limit, preserving excess rows for subsequent pulls.
+    pub fn pull_partial_data(&self, limit: &DataBlockLimit) -> Option<Result<DataBlock>> {
+        unsafe {
+            UpdateTrigger::update_input(&self.update_trigger);
+            // retrieve data without changing flags to prevent concurrent modifications
+            match self.shared.swap(std::ptr::null_mut(), 0, 0) {
                 address if address.is_null() => None,
                 address => {
-                    let block = (*Box::from_raw(address)).0;
-                    if let Ok(data_block) = block.as_ref() {
-                        ExecutorStats::record_thread_tracker(data_block.num_rows());
+                    let data_block = (*Box::from_raw(address)).0;
+                    if let Ok(data_block) = &data_block {
+                        let total_rows = data_block.num_rows();
+                        let total_bytes = data_block.memory_size();
+                        let rows_to_take = limit.rows_to_take(total_rows, total_bytes);
+                        // dbg!(rows_to_take, total_rows);
+                        // check if we need to split based on either rows or bytes limit
+                        if rows_to_take < total_rows {
+                            let (need, remain) = data_block.split_at(rows_to_take);
+                            let remain_data = Box::into_raw(Box::new(SharedData(Ok(remain))));
+                            // put back the remainder without changing flags
+                            // this keeps HAS_DATA set so we will have next pull to retrieve it
+                            self.shared.swap(remain_data, 0, 0);
+                            return Some(Ok(need));
+                        }
                     }
-                    Some(block)
+                    // take entire datablock and unset the flags to signal ready for next loop
+                    self.shared.set_flags(0, HAS_DATA | NEED_DATA);
+                    Some(data_block)
                 }
             }
         }