Skip to content

Commit 797019c

Browse files
committed
[native] Refactor TaskManagerTest to not create common::WriterOptions instance
After velox refactoring, we should not create dwio::common::WriterOptions instance anymore. Update TaskManagerTest.cpp
1 parent 67e423a commit 797019c

File tree

1 file changed

+7
-3
lines changed

1 file changed

+7
-3
lines changed

presto-native-execution/presto_cpp/main/tests/TaskManagerTest.cpp

Lines changed: 7 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -29,6 +29,7 @@
2929
#include "velox/dwio/common/FileSink.h"
3030
#include "velox/dwio/common/WriterFactory.h"
3131
#include "velox/dwio/common/tests/utils/BatchMaker.h"
32+
#include "velox/dwio/dwrf/writer/Writer.h"
3233
#include "velox/exec/Exchange.h"
3334
#include "velox/exec/Values.h"
3435
#include "velox/exec/tests/utils/PlanBuilder.h"
@@ -269,6 +270,8 @@ class TaskManagerTest : public testing::Test {
269270
"http://{}:{}",
270271
serverAddress.getAddressStr(),
271272
serverAddress.getPort()));
273+
writerFactory_ =
274+
dwio::common::getWriterFactory(dwio::common::FileFormat::DWRF);
272275
}
273276

274277
void TearDown() override {
@@ -298,13 +301,13 @@ class TaskManagerTest : public testing::Test {
298301
void writeToFile(
299302
const std::string& filePath,
300303
const std::vector<RowVectorPtr>& vectors) {
301-
auto options = std::make_shared<dwio::common::WriterOptions>();
304+
auto options = writerFactory_->createWriterOptions();
302305
options->schema = rowType_;
303306
options->memoryPool = rootPool_.get();
304307
auto sink = std::make_unique<dwio::common::LocalFileSink>(
305308
filePath, dwio::common::FileSink::Options{});
306-
auto writer = dwio::common::getWriterFactory(dwio::common::FileFormat::DWRF)
307-
->createWriter(std::move(sink), options);
309+
auto writer =
310+
writerFactory_->createWriter(std::move(sink), std::move(options));
308311

309312
for (size_t i = 0; i < vectors.size(); ++i) {
310313
writer->write(vectors[i]);
@@ -673,6 +676,7 @@ class TaskManagerTest : public testing::Test {
673676
long splitSequenceId_{0};
674677
std::shared_ptr<http::HttpClientConnectionPool> connPool_ =
675678
std::make_shared<http::HttpClientConnectionPool>();
679+
std::shared_ptr<dwio::common::WriterFactory> writerFactory_;
676680
};
677681

678682
// Runs "select * from t where c0 % 5 = 0" query.

0 commit comments

Comments
 (0)