From 0f0e6c5ee43ab755732ad34c0deaf3aa0a657cae Mon Sep 17 00:00:00 2001
From: Chigozie Joshua <36326251+guzzit@users.noreply.github.com>
Date: Sun, 16 Oct 2022 13:47:40 +0100
Subject: [PATCH] feat(batch): setup benchmark for LimitExecutor (#5589)

* added limit executor benchmark

* fixed clippy error

Co-authored-by: Bugen Zhao <i@bugenzhao.com>
---
 src/batch/Cargo.toml            |  4 +++
 src/batch/benches/limit.rs      | 64 +++++++++++++++++++++++++++++++++
 src/batch/src/executor/limit.rs | 17 +++++++--
 3 files changed, 82 insertions(+), 3 deletions(-)
 create mode 100644 src/batch/benches/limit.rs

diff --git a/src/batch/Cargo.toml b/src/batch/Cargo.toml
index cc42204dcca35..7e83f946356eb 100644
--- a/src/batch/Cargo.toml
+++ b/src/batch/Cargo.toml
@@ -99,3 +99,7 @@ harness = false
 [[bench]]
 name = "expand"
 harness = false
+
+[[bench]]
+name = "limit"
+harness = false
diff --git a/src/batch/benches/limit.rs b/src/batch/benches/limit.rs
new file mode 100644
index 0000000000000..a152ba08eddc0
--- /dev/null
+++ b/src/batch/benches/limit.rs
@@ -0,0 +1,64 @@
+// Copyright 2022 Singularity Data
+//
+// 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.
+
+pub mod utils;
+
+use criterion::{criterion_group, criterion_main, BatchSize, BenchmarkId, Criterion};
+use risingwave_batch::executor::{BoxedExecutor, LimitExecutor};
+use risingwave_common::types::DataType;
+use tikv_jemallocator::Jemalloc;
+use tokio::runtime::Runtime;
+use utils::{create_input, execute_executor};
+
+#[global_allocator]
+static GLOBAL: Jemalloc = Jemalloc;
+
+fn create_limit_executor(
+    chunk_size: usize,
+    chunk_num: usize,
+    offset: usize,
+    limit: usize,
+) -> BoxedExecutor {
+    let input = create_input(&[DataType::Int64], chunk_size, chunk_num);
+
+    Box::new(LimitExecutor::new(
+        input,
+        offset,
+        limit,
+        "LimitExecutor".into(),
+    ))
+}
+
+fn bench_limit(c: &mut Criterion) {
+    const SIZE: usize = 1024 * 1024;
+    let rt = Runtime::new().unwrap();
+
+    for chunk_size in &[32, 128, 512, 1024, 2048, 4096] {
+        c.bench_with_input(
+            BenchmarkId::new("LimitExecutor", chunk_size),
+            chunk_size,
+            |b, &chunk_size| {
+                let chunk_num = SIZE / chunk_size;
+                b.to_async(&rt).iter_batched(
+                    || create_limit_executor(chunk_size, chunk_num, 128, 128),
+                    |e| execute_executor(e),
+                    BatchSize::SmallInput,
+                );
+            },
+        );
+    }
+}
+
+criterion_group!(benches, bench_limit);
+criterion_main!(benches);
diff --git a/src/batch/src/executor/limit.rs b/src/batch/src/executor/limit.rs
index 5a42da294ad5f..7644ba361857d 100644
--- a/src/batch/src/executor/limit.rs
+++ b/src/batch/src/executor/limit.rs
@@ -51,12 +51,12 @@ impl BoxedExecutorBuilder for LimitExecutor {
         let limit = limit_node.get_limit() as usize;
         let offset = limit_node.get_offset() as usize;
 
-        Ok(Box::new(Self {
+        Ok(Box::new(Self::new(
             child,
             limit,
             offset,
-            identity: source.plan_node().get_identity().clone(),
-        }))
+            source.plan_node().get_identity().clone(),
+        )))
     }
 }
 
@@ -129,6 +129,17 @@ impl Executor for LimitExecutor {
     }
 }
 
+impl LimitExecutor {
+    pub fn new(child: BoxedExecutor, limit: usize, offset: usize, identity: String) -> Self {
+        Self {
+            child,
+            limit,
+            offset,
+            identity,
+        }
+    }
+}
+
 #[cfg(test)]
 mod tests {