|
29 | 29 | #include "velox/dwio/common/FileSink.h" |
30 | 30 | #include "velox/dwio/common/WriterFactory.h" |
31 | 31 | #include "velox/dwio/common/tests/utils/BatchMaker.h" |
| 32 | +#include "velox/dwio/dwrf/writer/Writer.h" |
32 | 33 | #include "velox/exec/Exchange.h" |
33 | 34 | #include "velox/exec/Values.h" |
34 | 35 | #include "velox/exec/tests/utils/PlanBuilder.h" |
@@ -269,6 +270,8 @@ class TaskManagerTest : public testing::Test { |
269 | 270 | "http://{}:{}", |
270 | 271 | serverAddress.getAddressStr(), |
271 | 272 | serverAddress.getPort())); |
| 273 | + writerFactory_ = |
| 274 | + dwio::common::getWriterFactory(dwio::common::FileFormat::DWRF); |
272 | 275 | } |
273 | 276 |
|
274 | 277 | void TearDown() override { |
@@ -298,13 +301,13 @@ class TaskManagerTest : public testing::Test { |
298 | 301 | void writeToFile( |
299 | 302 | const std::string& filePath, |
300 | 303 | const std::vector<RowVectorPtr>& vectors) { |
301 | | - auto options = std::make_shared<dwio::common::WriterOptions>(); |
| 304 | + auto options = writerFactory_->createWriterOptions(); |
302 | 305 | options->schema = rowType_; |
303 | 306 | options->memoryPool = rootPool_.get(); |
304 | 307 | auto sink = std::make_unique<dwio::common::LocalFileSink>( |
305 | 308 | 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)); |
308 | 311 |
|
309 | 312 | for (size_t i = 0; i < vectors.size(); ++i) { |
310 | 313 | writer->write(vectors[i]); |
@@ -673,6 +676,7 @@ class TaskManagerTest : public testing::Test { |
673 | 676 | long splitSequenceId_{0}; |
674 | 677 | std::shared_ptr<http::HttpClientConnectionPool> connPool_ = |
675 | 678 | std::make_shared<http::HttpClientConnectionPool>(); |
| 679 | + std::shared_ptr<dwio::common::WriterFactory> writerFactory_; |
676 | 680 | }; |
677 | 681 |
|
678 | 682 | // Runs "select * from t where c0 % 5 = 0" query. |
|
0 commit comments