[OPPRO-173] Make batch size configurable (facebookincubator#32)

rui-mo · Aug 18, 2022 · 16ab7dd · 16ab7dd
1 parent 5427c06
commit 16ab7dd
Show file tree

Hide file tree

Showing 2 changed files with 9 additions and 1 deletion.
diff --git 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
@@ -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_;