From 726fd3de5e006ef9a8c3fb67dce641fed66abf49 Mon Sep 17 00:00:00 2001
From: Rui Mo <rui.mo@intel.com>
Date: Thu, 28 Jul 2022 22:38:25 +0800
Subject: [PATCH] [OPPRO-173] Make batch size configurable (#32)

---
 velox/exec/TableScan.cpp | 8 +++++++-
 velox/exec/TableScan.h   | 2 ++
 2 files changed, 9 insertions(+), 1 deletion(-)

diff --git a/velox/exec/TableScan.cpp b/velox/exec/TableScan.cpp
index 661a92d965e2..b8f47f84230d 100644
--- a/velox/exec/TableScan.cpp
+++ b/velox/exec/TableScan.cpp
@@ -32,7 +32,8 @@ TableScan::TableScan(
           "TableScan"),
       tableHandle_(tableScanNode->tableHandle()),
       columnHandles_(tableScanNode->assignments()),
-      driverCtx_(driverCtx) {
+      driverCtx_(driverCtx),
+      preferredBatchSize_(driverCtx->queryConfig().preferredOutputBatchSize()) {
   connector_ = connector::getConnector(tableHandle_->connectorId());
 }
 
@@ -138,6 +139,11 @@ bool TableScan::isFinished() {
 }
 
 void TableScan::setBatchSize() {
+  if (preferredBatchSize_ != 1024) {
+    // Not the default value.
+    readBatchSize_ = preferredBatchSize_;
+    return;
+  }
   constexpr int64_t kMB = 1 << 20;
   auto estimate = dataSource_->estimatedRowSize();
   if (estimate == connector::DataSource::kUnknownRowSize) {
diff --git a/velox/exec/TableScan.h b/velox/exec/TableScan.h
index 8bba6d136297..1b4b69e345a7 100644
--- a/velox/exec/TableScan.h
+++ b/velox/exec/TableScan.h
@@ -69,6 +69,8 @@ class TableScan : public SourceOperator {
   std::unordered_map<column_index_t, std::shared_ptr<common::Filter>>
       pendingDynamicFilters_;
   int32_t readBatchSize_{kDefaultBatchSize};
+  // A preferred batch size from configuration.
+  uint32_t preferredBatchSize_;
 
   // String shown in ExceptionContext inside DataSource and LazyVector loading.
   std::string debugString_;