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> 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_;