From 3907faada684379d6078901d127de4859ceb0aae Mon Sep 17 00:00:00 2001 From: zouyunhe Date: Mon, 6 Jul 2026 06:20:46 +0000 Subject: [PATCH 1/6] support fs snapshot & restore states --- velox/connectors/Connector.h | 4 + .../filesystem/FileSystemDataSink.cpp | 152 +++++++++++++++++- .../filesystem/FileSystemDataSink.h | 31 +++- .../tests/FileSystemConnectorTest.cpp | 56 +++++++ velox/exec/TableWriter.h | 6 + 5 files changed, 239 insertions(+), 10 deletions(-) diff --git a/velox/connectors/Connector.h b/velox/connectors/Connector.h index 37ee5d02505..96720b6249c 100644 --- a/velox/connectors/Connector.h +++ b/velox/connectors/Connector.h @@ -242,6 +242,10 @@ class DataSink { /// Updates the current event-time watermark in milliseconds. virtual void setWatermark(int64_t watermark) {} + + /// Restores sink state from checkpoint records produced by snapshot(). + virtual void restoreState( + const std::vector& /*checkpointRecords*/) {} }; class DataSource { diff --git a/velox/connectors/filesystem/FileSystemDataSink.cpp b/velox/connectors/filesystem/FileSystemDataSink.cpp index bfa19fc6644..441a202c4af 100644 --- a/velox/connectors/filesystem/FileSystemDataSink.cpp +++ b/velox/connectors/filesystem/FileSystemDataSink.cpp @@ -17,6 +17,7 @@ #include "velox/connectors/filesystem/FileSystemDataSink.h" #include #include +#include #include "boost/uuid/uuid.hpp" #include "boost/uuid/uuid_generators.hpp" #include "boost/uuid/uuid_io.hpp" @@ -148,8 +149,8 @@ const std::unique_ptr FileSystemDataSink::createWriter( } const std::pair -FileSystemDataSink::getWriterFileNames() const { - return fileNameGenerator_->gen(); +FileSystemDataSink::getWriterFileNames(const std::string& bucketId) const { + return fileNameGenerator_->gen(bucketId); } namespace { @@ -181,9 +182,15 @@ std::string makePartitionDirectory( } } // namespace +std::string FileSystemDataSink::bucketIdForPartition( + const std::optional& partition) const { + return partition.value_or(""); +} + FsWriterParameters FileSystemDataSink::getWriterParameters( const std::optional& partition) const { - auto [targetFileName, writeFileName] = getWriterFileNames(); + auto [targetFileName, writeFileName] = + getWriterFileNames(bucketIdForPartition(partition)); return FsWriterParameters{ partition, targetFileName, @@ -554,7 +561,124 @@ std::vector FileSystemDataSink::snapshot(int64_t checkpointId) { pendingWriterInfo_.end()); pendingWriterInfo_.clear(); } - return {}; + return snapshotBucketState(checkpointId); +} + +std::vector FileSystemDataSink::snapshotBucketState( + int64_t checkpointId) const { + folly::dynamic state = folly::dynamic::object; + state["connector"] = "filesystem"; + state["checkpointId"] = checkpointId; + state["buckets"] = folly::dynamic::array; + state["pendingFiles"] = folly::dynamic::array; + + std::unordered_map counters = + fileNameGenerator_->partCounters(); + for (const auto& [_, writerInfos] : checkpointWriterInfo_) { + for (const auto& writerInfo : writerInfos) { + const auto& writerParams = writerInfo->writerParameters; + const auto bucketId = bucketIdForPartition(writerParams.partitionName()); + counters.try_emplace(bucketId, fileNameGenerator_->partCounter(bucketId)); + } + } + for (const auto& writerInfo : pendingWriterInfo_) { + const auto& writerParams = writerInfo->writerParameters; + const auto bucketId = bucketIdForPartition(writerParams.partitionName()); + counters.try_emplace(bucketId, fileNameGenerator_->partCounter(bucketId)); + } + + for (const auto& [bucketId, partCounter] : counters) { + folly::dynamic bucket = folly::dynamic::object; + bucket["bucketId"] = bucketId; + bucket["partCounter"] = partCounter; + state["buckets"].push_back(std::move(bucket)); + } + + for (const auto& [pendingCheckpointId, writerInfos] : checkpointWriterInfo_) { + for (const auto& writerInfo : writerInfos) { + const auto& writerParams = writerInfo->writerParameters; + folly::dynamic pendingFile = folly::dynamic::object; + pendingFile["checkpointId"] = pendingCheckpointId; + pendingFile["bucketId"] = + bucketIdForPartition(writerParams.partitionName()); + pendingFile["partitionName"] = writerParams.partitionName().value_or(""); + pendingFile["hasPartition"] = writerParams.partitionName().has_value(); + pendingFile["targetFileName"] = writerParams.targetFileName(); + pendingFile["targetDirectory"] = writerParams.targetDirectory(); + pendingFile["writeFileName"] = writerParams.writeFileName(); + pendingFile["writeDirectory"] = writerParams.writeDirectory(); + pendingFile["createTime"] = + static_cast(writerInfo->createTime()); + pendingFile["committed"] = writerInfo->isCommitted(); + pendingFile["fileRolled"] = writerInfo->isFileRolled(); + state["pendingFiles"].push_back(std::move(pendingFile)); + } + } + + return {folly::toJson(state)}; +} + +void FileSystemDataSink::restoreState( + const std::vector& checkpointRecords) { + for (const auto& checkpointRecord : checkpointRecords) { + folly::dynamic state; + try { + state = folly::parseJson(checkpointRecord); + } catch (const std::exception&) { + continue; + } + if (!state.isObject() || !state.count("connector") || + state["connector"].asString() != "filesystem") { + continue; + } + + if (state.count("buckets")) { + for (const auto& bucket : state["buckets"]) { + const auto bucketId = bucket["bucketId"].asString(); + const auto partCounter = bucket["partCounter"].asInt(); + VELOX_CHECK_GE(partCounter, 0, "Invalid restored part counter"); + fileNameGenerator_->restorePartCounter( + bucketId, static_cast(partCounter)); + } + } + + if (!state.count("pendingFiles")) { + continue; + } + int32_t restoredFileIndex = 0; + for (const auto& pendingFile : state["pendingFiles"]) { + const auto pendingCheckpointId = pendingFile["checkpointId"].asInt(); + const bool hasPartition = pendingFile["hasPartition"].asBool(); + std::optional partitionName = std::nullopt; + if (hasPartition) { + partitionName = pendingFile["partitionName"].asString(); + partitionCreateTimeMs_.try_emplace( + partitionName.value(), pendingFile["createTime"].asInt() * 1000); + } + + auto* connectorPool = queryCtx_->connectorMemoryPool(); + auto writerPool = connectorPool->addAggregateChild(fmt::format( + "{}.restored-file-system-sink.{}.{}", + connectorPool->name(), + pendingCheckpointId, + restoredFileIndex++)); + auto sinkPool = createSinkPool(writerPool); + auto writerInfo = std::make_shared( + FsWriterParameters{ + partitionName, + pendingFile["targetFileName"].asString(), + pendingFile["targetDirectory"].asString(), + pendingFile["writeFileName"].asString(), + pendingFile["writeDirectory"].asString()}, + std::move(writerPool), + std::move(sinkPool), + static_cast(pendingFile["createTime"].asInt())); + writerInfo->setCommitted(pendingFile["committed"].asBool()); + writerInfo->setFileRolled(pendingFile["fileRolled"].asBool()); + checkpointWriterInfo_[pendingCheckpointId].emplace_back( + std::move(writerInfo)); + } + } } // Renames in-progress files and returns paths ready for Flink-side partition @@ -762,11 +886,13 @@ const int64_t FileSystemDataSink::extractTimestampFromPartitionName( return timestamp.getSeconds(); } -const std::pair FsFileNameGenerator::gen() const { +const std::pair FsFileNameGenerator::gen( + const std::string& bucketId) const { std::pair fileNames; std::stringstream targetFileName; std::stringstream writeFileName; - targetFileName << "part-" << prefix_ << "-" << taskId_ << "-" << partCounter_ + auto& partCounter = partCounters_[bucketId]; + targetFileName << "part-" << prefix_ << "-" << taskId_ << "-" << partCounter << suffix_; boost::uuids::random_generator generator; boost::uuids::uuid uuid = generator(); @@ -774,8 +900,20 @@ const std::pair FsFileNameGenerator::gen() const { << to_string(uuid); fileNames.first = targetFileName.str(); fileNames.second = writeFileName.str(); - partCounter_++; + partCounter++; return fileNames; } +uint64_t FsFileNameGenerator::partCounter(const std::string& bucketId) const { + const auto it = partCounters_.find(bucketId); + return it == partCounters_.end() ? 0 : it->second; +} + +void FsFileNameGenerator::restorePartCounter( + const std::string& bucketId, + uint64_t partCounter) const { + auto& currentCounter = partCounters_[bucketId]; + currentCounter = std::max(currentCounter, partCounter); +} + } // namespace facebook::velox::connector::filesystem diff --git a/velox/connectors/filesystem/FileSystemDataSink.h b/velox/connectors/filesystem/FileSystemDataSink.h index 157a0897dd3..c308250ae11 100644 --- a/velox/connectors/filesystem/FileSystemDataSink.h +++ b/velox/connectors/filesystem/FileSystemDataSink.h @@ -18,6 +18,8 @@ #include #include #include +#include +#include #include "boost/algorithm/string.hpp" #include "folly/container/F14Map.h" #include "velox/common/file/FileSystems.h" @@ -94,6 +96,10 @@ struct FsWriterInfo : public hive::HiveWriterInfo { fileRolled_ = fileRolled; } + time_t createTime() const { + return createTime_; + } + private: const time_t createTime_; bool committed_ = false; @@ -110,13 +116,23 @@ class FsFileNameGenerator { const std::string& taskId) : prefix_(prefix), suffix_(suffix), taskId_(taskId) {} - const std::pair gen() const; + const std::pair gen( + const std::string& bucketId) const; + + uint64_t partCounter(const std::string& bucketId) const; + + const std::unordered_map& partCounters() const { + return partCounters_; + } + + void restorePartCounter(const std::string& bucketId, uint64_t partCounter) + const; private: const std::string prefix_; const std::string suffix_; const std::string taskId_; - mutable uint64_t partCounter_ = 0; + mutable std::unordered_map partCounters_; }; class FsPartitionIdGenerator : public hive::PartitionIdGenerator { @@ -224,6 +240,9 @@ class FileSystemDataSink : public DataSink { /// Receives event-time watermark in milliseconds. void setWatermark(int64_t watermark) override; + /// Restores bucket counters and pending files from checkpoint records. + void restoreState(const std::vector& checkpointRecords) override; + // For test. const std::vector& getWriteInfos() { return writerInfo_; @@ -326,11 +345,17 @@ class FileSystemDataSink : public DataSink { // prefix to the target file for write file name. The coordinator (or driver // for Presto on spark) will rename the write file to target file to commit // the table write when update the metadata store. - const std::pair getWriterFileNames() const; + const std::pair getWriterFileNames( + const std::string& bucketId) const; FsWriterParameters getWriterParameters( const std::optional& partition) const; + std::string bucketIdForPartition( + const std::optional& partition) const; + + std::vector snapshotBucketState(int64_t checkpointId) const; + // Get the HiveWriter corresponding to the row from partitionIds. FOLLY_ALWAYS_INLINE FsWriterId getWriterId(size_t row) const; diff --git a/velox/connectors/filesystem/tests/FileSystemConnectorTest.cpp b/velox/connectors/filesystem/tests/FileSystemConnectorTest.cpp index 7828f0333e5..a20587f9a6c 100644 --- a/velox/connectors/filesystem/tests/FileSystemConnectorTest.cpp +++ b/velox/connectors/filesystem/tests/FileSystemConnectorTest.cpp @@ -333,6 +333,62 @@ TEST_F( fs->remove(postBarrierTargetPath); } +TEST_F(FileSystemConnectorTest, testRestoreBucketStateKeepsPartCounter) { + std::shared_ptr fsConnector = + getConnector(fileSystemConnectorId); + auto queryCtx = createQueryCtx(); + std::unique_ptr dataSink = fsConnector->createDataSink( + inputType, + createFsTableHandle({}, {}), + queryCtx.get(), + CommitStrategy::kNoCommit); + auto* fsDataSink = reinterpret_cast(dataSink.get()); + ASSERT_TRUE(fsDataSink != nullptr); + + fsDataSink->appendData(createSingleInputRow()); + const auto checkpointTargetPath = + fsDataSink->getWriteInfos()[0]->writerParameters.targetDirectory() + "/" + + fsDataSink->getWriteInfos()[0]->writerParameters.targetFileName(); + const std::vector checkpointRecords = fsDataSink->snapshot(10); + ASSERT_EQ(checkpointRecords.size(), 1); + dataSink.reset(); + + auto restoredQueryCtx = createQueryCtx(); + std::unique_ptr restoredDataSink = fsConnector->createDataSink( + inputType, + createFsTableHandle({}, {}), + restoredQueryCtx.get(), + CommitStrategy::kNoCommit); + auto* restoredFsDataSink = + reinterpret_cast(restoredDataSink.get()); + ASSERT_TRUE(restoredFsDataSink != nullptr); + restoredFsDataSink->restoreState(checkpointRecords); + + restoredFsDataSink->appendData(createSingleInputRow()); + const auto restoredTargetPath = + restoredFsDataSink->getWriteInfos()[0] + ->writerParameters.targetDirectory() + + "/" + restoredFsDataSink->getWriteInfos()[0] + ->writerParameters.targetFileName(); + ASSERT_NE(checkpointTargetPath, restoredTargetPath); + + const std::vector committed = restoredFsDataSink->commit(10); + ASSERT_EQ(committed.size(), 1); + ASSERT_EQ(committed[0], checkpointTargetPath); + + auto fs = + filesystems::getFileSystem(dataPath, fsConnector->connectorConfig()); + ASSERT_TRUE(fs->exists(checkpointTargetPath)); + fs->remove(checkpointTargetPath); + + restoredFsDataSink->snapshot(11); + const std::vector committed2 = restoredFsDataSink->commit(11); + ASSERT_EQ(committed2.size(), 1); + ASSERT_EQ(committed2[0], restoredTargetPath); + ASSERT_TRUE(fs->exists(restoredTargetPath)); + fs->remove(restoredTargetPath); +} + TEST_F(FileSystemConnectorTest, testNonPartitionedFileCommit) { std::shared_ptr fsConnector = getConnector(fileSystemConnectorId); diff --git a/velox/exec/TableWriter.h b/velox/exec/TableWriter.h index 9fa8664d072..7a6c3841ce0 100644 --- a/velox/exec/TableWriter.h +++ b/velox/exec/TableWriter.h @@ -132,6 +132,12 @@ class TableWriter : public Operator { } } + void restoreState(const std::vector& checkpointRecords) override { + if (dataSink_) { + dataSink_->restoreState(checkpointRecords); + } + } + virtual bool needsInput() const override { return true; } From 429ab89d802e747cb20795916b03860c4c46f3fd Mon Sep 17 00:00:00 2001 From: zouyunhe Date: Wed, 8 Jul 2026 03:05:20 +0000 Subject: [PATCH 2/6] fix --- .../filesystem/FileSystemDataSink.cpp | 21 ++++++++++++++----- 1 file changed, 16 insertions(+), 5 deletions(-) diff --git a/velox/connectors/filesystem/FileSystemDataSink.cpp b/velox/connectors/filesystem/FileSystemDataSink.cpp index 441a202c4af..8127ca3d007 100644 --- a/velox/connectors/filesystem/FileSystemDataSink.cpp +++ b/velox/connectors/filesystem/FileSystemDataSink.cpp @@ -180,6 +180,7 @@ std::string makePartitionDirectory( } return (fs::path(tableDirectory) / partitionSubdirectory.value()).string(); } + } // namespace std::string FileSystemDataSink::bucketIdForPartition( @@ -741,11 +742,21 @@ std::vector FileSystemDataSink::commit(int64_t checkpointId) { fs_->rename(writeFileName, targetFileName); writerInfo->setFileRolled(true); } catch (const std::exception& e) { - VELOX_FAIL( - "Failed to rename file {} to target {}, exception: {}", - writeFileName, - targetFileName, - e.what()); + const auto targetExists = fs_->exists(targetFileName); + const auto writeExists = fs_->exists(writeFileName); + if (targetExists && !writeExists) { + LOG(INFO) << "Treat rename as already completed for file " + << writeFileName << " to target " << targetFileName + << ", targetExists=" << targetExists + << ", exception: " << e.what(); + writerInfo->setFileRolled(true); + } else { + VELOX_FAIL( + "Failed to rename file {} to target {}, exception: {}", + writeFileName, + targetFileName, + e.what()); + } } } From e4690fa2e84f8e6496d3ba6047368c6ecd6f195e Mon Sep 17 00:00:00 2001 From: zouyunhe Date: Tue, 14 Jul 2026 06:57:15 +0000 Subject: [PATCH 3/6] align with flink --- .../filesystem/FileSystemDataSink.cpp | 66 ++++++++----------- .../filesystem/FileSystemDataSink.h | 21 ++++-- .../tests/FileSystemConnectorTest.cpp | 34 ++++++++++ 3 files changed, 76 insertions(+), 45 deletions(-) diff --git a/velox/connectors/filesystem/FileSystemDataSink.cpp b/velox/connectors/filesystem/FileSystemDataSink.cpp index 8127ca3d007..9abf66dcf5f 100644 --- a/velox/connectors/filesystem/FileSystemDataSink.cpp +++ b/velox/connectors/filesystem/FileSystemDataSink.cpp @@ -570,31 +570,10 @@ std::vector FileSystemDataSink::snapshotBucketState( folly::dynamic state = folly::dynamic::object; state["connector"] = "filesystem"; state["checkpointId"] = checkpointId; - state["buckets"] = folly::dynamic::array; + // StreamingFileSink stores maxPartCounter in a separate partCounterState. + state["partCounter"] = fileNameGenerator_->maxPartCounter(); state["pendingFiles"] = folly::dynamic::array; - std::unordered_map counters = - fileNameGenerator_->partCounters(); - for (const auto& [_, writerInfos] : checkpointWriterInfo_) { - for (const auto& writerInfo : writerInfos) { - const auto& writerParams = writerInfo->writerParameters; - const auto bucketId = bucketIdForPartition(writerParams.partitionName()); - counters.try_emplace(bucketId, fileNameGenerator_->partCounter(bucketId)); - } - } - for (const auto& writerInfo : pendingWriterInfo_) { - const auto& writerParams = writerInfo->writerParameters; - const auto bucketId = bucketIdForPartition(writerParams.partitionName()); - counters.try_emplace(bucketId, fileNameGenerator_->partCounter(bucketId)); - } - - for (const auto& [bucketId, partCounter] : counters) { - folly::dynamic bucket = folly::dynamic::object; - bucket["bucketId"] = bucketId; - bucket["partCounter"] = partCounter; - state["buckets"].push_back(std::move(bucket)); - } - for (const auto& [pendingCheckpointId, writerInfos] : checkpointWriterInfo_) { for (const auto& writerInfo : writerInfos) { const auto& writerParams = writerInfo->writerParameters; @@ -633,13 +612,19 @@ void FileSystemDataSink::restoreState( continue; } - if (state.count("buckets")) { + // Prefer StreamingFileSink-style global max partCounter. Fall back to the + // legacy per-bucket buckets[].partCounter format by taking the max. + if (state.count("partCounter")) { + const auto partCounter = state["partCounter"].asInt(); + VELOX_CHECK_GE(partCounter, 0, "Invalid restored part counter"); + fileNameGenerator_->restoreMaxPartCounter( + static_cast(partCounter)); + } else if (state.count("buckets")) { for (const auto& bucket : state["buckets"]) { - const auto bucketId = bucket["bucketId"].asString(); const auto partCounter = bucket["partCounter"].asInt(); VELOX_CHECK_GE(partCounter, 0, "Invalid restored part counter"); - fileNameGenerator_->restorePartCounter( - bucketId, static_cast(partCounter)); + fileNameGenerator_->restoreMaxPartCounter( + static_cast(partCounter)); } } @@ -899,10 +884,15 @@ const int64_t FileSystemDataSink::extractTimestampFromPartitionName( const std::pair FsFileNameGenerator::gen( const std::string& bucketId) const { + // Initialize this bucket's counter from the global max, matching + // StreamingFileSink Buckets#getOrCreateBucketForBucketId. + auto [it, inserted] = partCounters_.try_emplace(bucketId, maxPartCounter_); + (void)inserted; + auto& partCounter = it->second; + std::pair fileNames; std::stringstream targetFileName; std::stringstream writeFileName; - auto& partCounter = partCounters_[bucketId]; targetFileName << "part-" << prefix_ << "-" << taskId_ << "-" << partCounter << suffix_; boost::uuids::random_generator generator; @@ -912,19 +902,19 @@ const std::pair FsFileNameGenerator::gen( fileNames.first = targetFileName.str(); fileNames.second = writeFileName.str(); partCounter++; + // Keep the global max in sync so inactive buckets reopened later do not + // restart from 0 and overwrite committed parts. + maxPartCounter_ = std::max(maxPartCounter_, partCounter); return fileNames; } -uint64_t FsFileNameGenerator::partCounter(const std::string& bucketId) const { - const auto it = partCounters_.find(bucketId); - return it == partCounters_.end() ? 0 : it->second; -} - -void FsFileNameGenerator::restorePartCounter( - const std::string& bucketId, - uint64_t partCounter) const { - auto& currentCounter = partCounters_[bucketId]; - currentCounter = std::max(currentCounter, partCounter); +void FsFileNameGenerator::restoreMaxPartCounter(uint64_t partCounter) const { + maxPartCounter_ = std::max(maxPartCounter_, partCounter); + // Buckets created after restore start from the restored max, as in + // StreamingFileSink Buckets#initializePartCounter / restoreBucket. + for (auto& [_, counter] : partCounters_) { + counter = std::max(counter, maxPartCounter_); + } } } // namespace facebook::velox::connector::filesystem diff --git a/velox/connectors/filesystem/FileSystemDataSink.h b/velox/connectors/filesystem/FileSystemDataSink.h index c308250ae11..ef0386f0a4d 100644 --- a/velox/connectors/filesystem/FileSystemDataSink.h +++ b/velox/connectors/filesystem/FileSystemDataSink.h @@ -108,6 +108,9 @@ struct FsWriterInfo : public hive::HiveWriterInfo { using FsWriterInfoPtr = std::shared_ptr; +/// Generates part file names using a StreamingFileSink-style counter model: +/// each active bucket keeps an in-memory partCounter initialized from the +/// global maxPartCounter; only maxPartCounter is checkpointed. class FsFileNameGenerator { public: FsFileNameGenerator( @@ -119,19 +122,23 @@ class FsFileNameGenerator { const std::pair gen( const std::string& bucketId) const; - uint64_t partCounter(const std::string& bucketId) const; - - const std::unordered_map& partCounters() const { - return partCounters_; + /// Global max part counter across all buckets (checkpointed like Flink's + /// partCounterState). + uint64_t maxPartCounter() const { + return maxPartCounter_; } - void restorePartCounter(const std::string& bucketId, uint64_t partCounter) - const; + /// Restores the global max part counter. New/restored buckets start from + /// this value so previously committed part files are not overwritten. + void restoreMaxPartCounter(uint64_t partCounter) const; private: const std::string prefix_; const std::string suffix_; const std::string taskId_; + /// Max part counter used across all buckets for this sink task. + mutable uint64_t maxPartCounter_{0}; + /// Per-bucket counters for currently active buckets (not checkpointed). mutable std::unordered_map partCounters_; }; @@ -240,7 +247,7 @@ class FileSystemDataSink : public DataSink { /// Receives event-time watermark in milliseconds. void setWatermark(int64_t watermark) override; - /// Restores bucket counters and pending files from checkpoint records. + /// Restores max part counter and pending files from checkpoint records. void restoreState(const std::vector& checkpointRecords) override; // For test. diff --git a/velox/connectors/filesystem/tests/FileSystemConnectorTest.cpp b/velox/connectors/filesystem/tests/FileSystemConnectorTest.cpp index a20587f9a6c..4daa18df330 100644 --- a/velox/connectors/filesystem/tests/FileSystemConnectorTest.cpp +++ b/velox/connectors/filesystem/tests/FileSystemConnectorTest.cpp @@ -351,6 +351,9 @@ TEST_F(FileSystemConnectorTest, testRestoreBucketStateKeepsPartCounter) { fsDataSink->getWriteInfos()[0]->writerParameters.targetFileName(); const std::vector checkpointRecords = fsDataSink->snapshot(10); ASSERT_EQ(checkpointRecords.size(), 1); + // StreamingFileSink-style state: global max partCounter, not per-bucket map. + ASSERT_NE(checkpointRecords[0].find("\"partCounter\""), std::string::npos); + ASSERT_EQ(checkpointRecords[0].find("\"buckets\""), std::string::npos); dataSink.reset(); auto restoredQueryCtx = createQueryCtx(); @@ -389,6 +392,37 @@ TEST_F(FileSystemConnectorTest, testRestoreBucketStateKeepsPartCounter) { fs->remove(restoredTargetPath); } +TEST_F(FileSystemConnectorTest, testRestoreLegacyPerBucketPartCounterState) { + std::shared_ptr fsConnector = + getConnector(fileSystemConnectorId); + auto restoredQueryCtx = createQueryCtx(); + std::unique_ptr restoredDataSink = fsConnector->createDataSink( + inputType, + createFsTableHandle({}, {}), + restoredQueryCtx.get(), + CommitStrategy::kNoCommit); + auto* restoredFsDataSink = + reinterpret_cast(restoredDataSink.get()); + ASSERT_TRUE(restoredFsDataSink != nullptr); + + // Legacy state format with per-bucket counters; restore should take the max. + const std::string legacyState = R"({ + "connector":"filesystem", + "checkpointId":1, + "buckets":[ + {"bucketId":"","partCounter":3}, + {"bucketId":"p=0","partCounter":5} + ], + "pendingFiles":[] + })"; + restoredFsDataSink->restoreState({legacyState}); + restoredFsDataSink->appendData(createSingleInputRow()); + const auto targetFileName = + restoredFsDataSink->getWriteInfos()[0]->writerParameters.targetFileName(); + // After restoring max=5, the next part file must use counter 5. + ASSERT_NE(targetFileName.find("-5"), std::string::npos); +} + TEST_F(FileSystemConnectorTest, testNonPartitionedFileCommit) { std::shared_ptr fsConnector = getConnector(fileSystemConnectorId); From 51131da76e5af9355ebb22e49e300b1a52be6ca4 Mon Sep 17 00:00:00 2001 From: zouyunhe Date: Tue, 14 Jul 2026 08:34:58 +0000 Subject: [PATCH 4/6] fix code format --- velox/connectors/filesystem/FileSystemDataSink.cpp | 2 +- .../filesystem/tests/FileSystemConnectorTest.cpp | 9 ++++----- velox/exec/TableWriter.h | 3 ++- 3 files changed, 7 insertions(+), 7 deletions(-) diff --git a/velox/connectors/filesystem/FileSystemDataSink.cpp b/velox/connectors/filesystem/FileSystemDataSink.cpp index 9abf66dcf5f..56dd2432378 100644 --- a/velox/connectors/filesystem/FileSystemDataSink.cpp +++ b/velox/connectors/filesystem/FileSystemDataSink.cpp @@ -16,8 +16,8 @@ #include "velox/connectors/filesystem/FileSystemDataSink.h" #include -#include #include +#include #include "boost/uuid/uuid.hpp" #include "boost/uuid/uuid_generators.hpp" #include "boost/uuid/uuid_io.hpp" diff --git a/velox/connectors/filesystem/tests/FileSystemConnectorTest.cpp b/velox/connectors/filesystem/tests/FileSystemConnectorTest.cpp index 4daa18df330..d2d06969fd3 100644 --- a/velox/connectors/filesystem/tests/FileSystemConnectorTest.cpp +++ b/velox/connectors/filesystem/tests/FileSystemConnectorTest.cpp @@ -368,11 +368,10 @@ TEST_F(FileSystemConnectorTest, testRestoreBucketStateKeepsPartCounter) { restoredFsDataSink->restoreState(checkpointRecords); restoredFsDataSink->appendData(createSingleInputRow()); - const auto restoredTargetPath = - restoredFsDataSink->getWriteInfos()[0] - ->writerParameters.targetDirectory() + - "/" + restoredFsDataSink->getWriteInfos()[0] - ->writerParameters.targetFileName(); + const auto restoredTargetPath = restoredFsDataSink->getWriteInfos()[0] + ->writerParameters.targetDirectory() + + "/" + + restoredFsDataSink->getWriteInfos()[0]->writerParameters.targetFileName(); ASSERT_NE(checkpointTargetPath, restoredTargetPath); const std::vector committed = restoredFsDataSink->commit(10); diff --git a/velox/exec/TableWriter.h b/velox/exec/TableWriter.h index 7a6c3841ce0..810364a53f7 100644 --- a/velox/exec/TableWriter.h +++ b/velox/exec/TableWriter.h @@ -132,7 +132,8 @@ class TableWriter : public Operator { } } - void restoreState(const std::vector& checkpointRecords) override { + void restoreState( + const std::vector& checkpointRecords) override { if (dataSink_) { dataSink_->restoreState(checkpointRecords); } From 37861610b211c0a3415b8b87f2a91ba611f011c2 Mon Sep 17 00:00:00 2001 From: zouyunhe Date: Tue, 14 Jul 2026 11:42:04 +0000 Subject: [PATCH 5/6] fix reviews --- .../filesystem/FileSystemDataSink.cpp | 79 +++++++++++++++---- .../filesystem/FileSystemDataSink.h | 22 +++--- 2 files changed, 76 insertions(+), 25 deletions(-) diff --git a/velox/connectors/filesystem/FileSystemDataSink.cpp b/velox/connectors/filesystem/FileSystemDataSink.cpp index 56dd2432378..24830b61615 100644 --- a/velox/connectors/filesystem/FileSystemDataSink.cpp +++ b/velox/connectors/filesystem/FileSystemDataSink.cpp @@ -91,15 +91,11 @@ FileSystemDataSink::FileSystemDataSink( dataChannels_( getNonPartitionChannels(partitionChannels_, inputType_->size())), writerFactory_(dwio::common::getWriterFactory(writeConfig_->getFormat())), - fileNameGenerator_( - std::make_shared( - queryCtx_->sessionProperties()->get( - "query_uuid", - ""), - "", - queryCtx_->sessionProperties()->get( - "task_index", - "0"))) {} + fileNameGenerator_(std::make_shared( + queryCtx_->sessionProperties()->get("query_uuid", ""), + "", + queryCtx_->sessionProperties()->get( + "task_index", "0"))) {} const std::unique_ptr FileSystemDataSink::createWriter( const std::string& writePath, @@ -149,7 +145,7 @@ const std::unique_ptr FileSystemDataSink::createWriter( } const std::pair -FileSystemDataSink::getWriterFileNames(const std::string& bucketId) const { +FileSystemDataSink::getWriterFileNames(const std::string& bucketId) { return fileNameGenerator_->gen(bucketId); } @@ -181,6 +177,43 @@ std::string makePartitionDirectory( return (fs::path(tableDirectory) / partitionSubdirectory.value()).string(); } +void validatePendingFileState(const folly::dynamic& pendingFile) { + VELOX_USER_CHECK( + pendingFile.isObject(), + "Invalid filesystem checkpoint state: pending file must be an object"); + const auto requireInt = [&](const char* key) { + VELOX_USER_CHECK( + pendingFile.count(key) && pendingFile[key].isInt(), + "Invalid filesystem checkpoint state: pending file missing int field " + "'{}'", + key); + }; + const auto requireBool = [&](const char* key) { + VELOX_USER_CHECK( + pendingFile.count(key) && pendingFile[key].isBool(), + "Invalid filesystem checkpoint state: pending file missing bool field " + "'{}'", + key); + }; + const auto requireString = [&](const char* key) { + VELOX_USER_CHECK( + pendingFile.count(key) && pendingFile[key].isString(), + "Invalid filesystem checkpoint state: pending file missing string " + "field '{}'", + key); + }; + requireInt("checkpointId"); + requireBool("hasPartition"); + requireString("partitionName"); + requireString("targetFileName"); + requireString("targetDirectory"); + requireString("writeFileName"); + requireString("writeDirectory"); + requireInt("createTime"); + requireBool("committed"); + requireBool("fileRolled"); +} + } // namespace std::string FileSystemDataSink::bucketIdForPartition( @@ -189,7 +222,7 @@ std::string FileSystemDataSink::bucketIdForPartition( } FsWriterParameters FileSystemDataSink::getWriterParameters( - const std::optional& partition) const { + const std::optional& partition) { auto [targetFileName, writeFileName] = getWriterFileNames(bucketIdForPartition(partition)); return FsWriterParameters{ @@ -608,6 +641,7 @@ void FileSystemDataSink::restoreState( continue; } if (!state.isObject() || !state.count("connector") || + !state["connector"].isString() || state["connector"].asString() != "filesystem") { continue; } @@ -615,14 +649,25 @@ void FileSystemDataSink::restoreState( // Prefer StreamingFileSink-style global max partCounter. Fall back to the // legacy per-bucket buckets[].partCounter format by taking the max. if (state.count("partCounter")) { + VELOX_USER_CHECK( + state["partCounter"].isInt(), + "Invalid filesystem checkpoint state: partCounter must be an int"); const auto partCounter = state["partCounter"].asInt(); - VELOX_CHECK_GE(partCounter, 0, "Invalid restored part counter"); + VELOX_USER_CHECK_GE(partCounter, 0, "Invalid restored part counter"); fileNameGenerator_->restoreMaxPartCounter( static_cast(partCounter)); } else if (state.count("buckets")) { + VELOX_USER_CHECK( + state["buckets"].isArray(), + "Invalid filesystem checkpoint state: buckets must be an array"); for (const auto& bucket : state["buckets"]) { + VELOX_USER_CHECK( + bucket.isObject() && bucket.count("partCounter") && + bucket["partCounter"].isInt(), + "Invalid filesystem checkpoint state: each bucket must be an " + "object with an int partCounter"); const auto partCounter = bucket["partCounter"].asInt(); - VELOX_CHECK_GE(partCounter, 0, "Invalid restored part counter"); + VELOX_USER_CHECK_GE(partCounter, 0, "Invalid restored part counter"); fileNameGenerator_->restoreMaxPartCounter( static_cast(partCounter)); } @@ -631,8 +676,12 @@ void FileSystemDataSink::restoreState( if (!state.count("pendingFiles")) { continue; } + VELOX_USER_CHECK( + state["pendingFiles"].isArray(), + "Invalid filesystem checkpoint state: pendingFiles must be an array"); int32_t restoredFileIndex = 0; for (const auto& pendingFile : state["pendingFiles"]) { + validatePendingFileState(pendingFile); const auto pendingCheckpointId = pendingFile["checkpointId"].asInt(); const bool hasPartition = pendingFile["hasPartition"].asBool(); std::optional partitionName = std::nullopt; @@ -883,7 +932,7 @@ const int64_t FileSystemDataSink::extractTimestampFromPartitionName( } const std::pair FsFileNameGenerator::gen( - const std::string& bucketId) const { + const std::string& bucketId) { // Initialize this bucket's counter from the global max, matching // StreamingFileSink Buckets#getOrCreateBucketForBucketId. auto [it, inserted] = partCounters_.try_emplace(bucketId, maxPartCounter_); @@ -908,7 +957,7 @@ const std::pair FsFileNameGenerator::gen( return fileNames; } -void FsFileNameGenerator::restoreMaxPartCounter(uint64_t partCounter) const { +void FsFileNameGenerator::restoreMaxPartCounter(uint64_t partCounter) { maxPartCounter_ = std::max(maxPartCounter_, partCounter); // Buckets created after restore start from the restored max, as in // StreamingFileSink Buckets#initializePartCounter / restoreBucket. diff --git a/velox/connectors/filesystem/FileSystemDataSink.h b/velox/connectors/filesystem/FileSystemDataSink.h index ef0386f0a4d..9bec4811c02 100644 --- a/velox/connectors/filesystem/FileSystemDataSink.h +++ b/velox/connectors/filesystem/FileSystemDataSink.h @@ -84,11 +84,11 @@ struct FsWriterInfo : public hive::HiveWriterInfo { committed_ = committed; } - bool isCommitted() { + bool isCommitted() const { return committed_; } - bool isFileRolled() { + bool isFileRolled() const { return fileRolled_; } @@ -119,8 +119,8 @@ class FsFileNameGenerator { const std::string& taskId) : prefix_(prefix), suffix_(suffix), taskId_(taskId) {} - const std::pair gen( - const std::string& bucketId) const; + // Stateful: advances per-bucket and global part counters. + const std::pair gen(const std::string& bucketId); /// Global max part counter across all buckets (checkpointed like Flink's /// partCounterState). @@ -130,16 +130,16 @@ class FsFileNameGenerator { /// Restores the global max part counter. New/restored buckets start from /// this value so previously committed part files are not overwritten. - void restoreMaxPartCounter(uint64_t partCounter) const; + void restoreMaxPartCounter(uint64_t partCounter); private: const std::string prefix_; const std::string suffix_; const std::string taskId_; /// Max part counter used across all buckets for this sink task. - mutable uint64_t maxPartCounter_{0}; + uint64_t maxPartCounter_{0}; /// Per-bucket counters for currently active buckets (not checkpointed). - mutable std::unordered_map partCounters_; + std::unordered_map partCounters_; }; class FsPartitionIdGenerator : public hive::PartitionIdGenerator { @@ -306,7 +306,8 @@ class FileSystemDataSink : public DataSink { std::vector rawPartitionRows_; std::vector partitionSizes_; const std::shared_ptr writerFactory_; - const std::shared_ptr fileNameGenerator_; + // Non-const: file name generation mutates part counters. + const std::shared_ptr fileNameGenerator_; std::shared_ptr fs_; int64_t watermark_{INT64_MIN}; // Partition creation time in milliseconds (system clock) for process-time @@ -352,11 +353,12 @@ class FileSystemDataSink : public DataSink { // prefix to the target file for write file name. The coordinator (or driver // for Presto on spark) will rename the write file to target file to commit // the table write when update the metadata store. + // Non-const: delegates to stateful FsFileNameGenerator::gen(). const std::pair getWriterFileNames( - const std::string& bucketId) const; + const std::string& bucketId); FsWriterParameters getWriterParameters( - const std::optional& partition) const; + const std::optional& partition); std::string bucketIdForPartition( const std::optional& partition) const; From 1da900093829f39a631be13eabe1acc8e13fe233 Mon Sep 17 00:00:00 2001 From: zouyunhe Date: Tue, 14 Jul 2026 12:23:55 +0000 Subject: [PATCH 6/6] fix code format --- velox/connectors/filesystem/FileSystemDataSink.cpp | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/velox/connectors/filesystem/FileSystemDataSink.cpp b/velox/connectors/filesystem/FileSystemDataSink.cpp index 24830b61615..77eefe46c36 100644 --- a/velox/connectors/filesystem/FileSystemDataSink.cpp +++ b/velox/connectors/filesystem/FileSystemDataSink.cpp @@ -95,7 +95,8 @@ FileSystemDataSink::FileSystemDataSink( queryCtx_->sessionProperties()->get("query_uuid", ""), "", queryCtx_->sessionProperties()->get( - "task_index", "0"))) {} + "task_index", + "0"))) {} const std::unique_ptr FileSystemDataSink::createWriter( const std::string& writePath,