fix: propagate json stats parquet close errors (#50870)

pr: #50856

Signed-off-by: luzhang <luzhang@zilliz.com>
Co-authored-by: luzhang <luzhang@zilliz.com>

internal/core/src/index/json_stats/JsonKeyStats.cpp

@@ -747,7 +747,10 @@ JsonKeyStats::BuildWithFieldData(const std::vector<FieldDataPtr>& field_datas) {
         ParquetWriterFactory::CreateContext(key_types_, remote_prefix);
     parquet_writer_->Init(std::move(writer_context));
     BuildKeyStats(field_datas);
-    parquet_writer_->Close();
+    auto close_status = parquet_writer_->Close();
+    AssertInfo(close_status.ok(),
+               "failed to close json stats parquet writer: {}",
+               close_status.ToString());
     bson_inverted_index_->BuildIndex();
 
     // write meta file with layout type map and other metadata

internal/core/src/index/json_stats/parquet_writer.cpp

@@ -19,9 +19,14 @@
 #include <arrow/array/array_primitive.h>
 #include <arrow/array/builder_binary.h>
 #include <arrow/array/builder_primitive.h>
+#include <exception>
 #include <arrow/io/file.h>
 #include <parquet/arrow/writer.h>
 #include <parquet/exception.h>
+#include <utility>
+
+#include "arrow/array/builder_base.h"
+#include "milvus-storage/packed/writer.h"
 
 namespace milvus::index {
 
@@ -41,22 +46,37 @@ JsonStatsParquetWriter::JsonStatsParquetWriter(
 }
 
 JsonStatsParquetWriter::~JsonStatsParquetWriter() {
-    Close();
-}
-
-void
-JsonStatsParquetWriter::Close() {
-    // check packed_writer_ initialized
-    if (packed_writer_) {
-        Flush();
-        // ignore close status here
-        auto _ = packed_writer_->Close();
+    try {
+        auto status = Close();
+        if (!status.ok()) {
+            LOG_WARN("failed to close json stats parquet writer: {}",
+                     status.ToString());
+        }
+    } catch (const std::exception& e) {
+        LOG_WARN("failed to close json stats parquet writer: {}", e.what());
     }
 }
 
-void
+arrow::Status
+JsonStatsParquetWriter::Close() {
+    if (!packed_writer_) {
+        return arrow::Status::OK();
+    }
+    auto flush_status = Flush();
+    if (!flush_status.ok()) {
+        return flush_status;
+    }
+    auto writer = std::move(packed_writer_);
+    return writer->Close();
+}
+
+arrow::Status
 JsonStatsParquetWriter::Flush() {
-    WriteCurrentBatch();
+    auto res = WriteCurrentBatch();
+    if (!res.ok()) {
+        return res.status();
+    }
+    return arrow::Status::OK();
 }
 
 void
@@ -74,7 +94,7 @@ JsonStatsParquetWriter::UpdatePathSizeMap(
     }
 }
 
-size_t
+arrow::Result<size_t>
 JsonStatsParquetWriter::WriteCurrentBatch() {
     if (unflushed_row_count_ == 0) {
         return 0;
@@ -85,8 +105,9 @@ JsonStatsParquetWriter::WriteCurrentBatch() {
     for (auto& builder : builders_) {
         std::shared_ptr<arrow::Array> array;
         auto status = builder->Finish(&array);
-        AssertInfo(
-            status.ok(), "failed to finish builder: {}", status.ToString());
+        if (!status.ok()) {
+            return status;
+        }
         arrays.push_back(std::move(array));
         builder->Reset();
     }
@@ -96,8 +117,10 @@ JsonStatsParquetWriter::WriteCurrentBatch() {
     auto batch = arrow::RecordBatch::Make(
         schema_, unflushed_row_count_, std::move(arrays));
     auto status = packed_writer_->Write(batch);
-    AssertInfo(
-        status.ok(), "failed to write batch, error: {}", status.ToString());
+    if (!status.ok()) {
+        unflushed_row_count_ = 0;
+        return status;
+    }
 
     auto res = unflushed_row_count_;
     unflushed_row_count_ = 0;
@@ -132,7 +155,10 @@ JsonStatsParquetWriter::AddCurrentRow() {
     unflushed_row_count_++;
     all_row_count_++;
     if (unflushed_row_count_ >= batch_size_) {
-        WriteCurrentBatch();
+        auto result = WriteCurrentBatch();
+        AssertInfo(result.ok(),
+                   "failed to write current json stats parquet batch: {}",
+                   result.status().ToString());
     }
     return all_row_count_;
 }
@@ -266,4 +292,4 @@ JsonStatsParquetWriter::AppendDataToBuilder(
     return ast;
 }
 
-}  // namespace milvus::index
+}  // namespace milvus::index

internal/core/src/index/json_stats/parquet_writer.h

@@ -189,13 +189,13 @@ class JsonStatsParquetWriter {
     void
     AppendSharedRow(const uint8_t* data, size_t length);
 
-    void
+    arrow::Status
     Flush();
 
-    void
+    arrow::Status
     Close();
 
-    size_t
+    arrow::Result<size_t>
     WriteCurrentBatch();
 
     size_t
@@ -243,4 +243,4 @@ class JsonStatsParquetWriter {
     size_t current_row_count_{0};
 };
 
-}  // namespace milvus::index
+}  // namespace milvus::index

internal/core/unittest/test_json_stats/test_parquet_writer.cpp

@@ -1,14 +1,91 @@
 #include <gtest/gtest.h>
-#include "index/json_stats/parquet_writer.h"
-#include <arrow/io/memory.h>
+#include <arrow/filesystem/localfs.h>
 #include <arrow/io/file.h>
+#include <arrow/io/interfaces.h>
+#include <arrow/io/memory.h>
+#include <arrow/result.h>
+#include <filesystem>
+#include <map>
 #include <parquet/arrow/writer.h>
 #include <memory>
+#include <set>
+#include <stdexcept>
 #include <string>
+#include <utility>
 #include <vector>
-#include <map>
+
+#include "index/json_stats/parquet_writer.h"
+#include "index/json_stats/utils.h"
+#include "storage/Util.h"
+#include "test_utils/Constants.h"
 
 namespace milvus::index {
+namespace {
+
+class CloseFailingOutputStream : public arrow::io::OutputStream {
+ public:
+    explicit CloseFailingOutputStream(
+        std::shared_ptr<arrow::io::OutputStream> delegate)
+        : delegate_(std::move(delegate)) {
+        set_mode(arrow::io::FileMode::WRITE);
+    }
+
+    arrow::Status
+    Close() override {
+        closed_ = true;
+        auto status = delegate_->Close();
+        if (!status.ok()) {
+            return status;
+        }
+        return arrow::Status::IOError("injected final close failure");
+    }
+
+    bool
+    closed() const override {
+        return closed_ || delegate_->closed();
+    }
+
+    arrow::Result<int64_t>
+    Tell() const override {
+        return delegate_->Tell();
+    }
+
+    arrow::Status
+    Write(const void* data, int64_t nbytes) override {
+        return delegate_->Write(data, nbytes);
+    }
+
+    arrow::Status
+    Flush() override {
+        return delegate_->Flush();
+    }
+
+    using arrow::io::Writable::Write;
+
+ private:
+    std::shared_ptr<arrow::io::OutputStream> delegate_;
+    bool closed_{false};
+};
+
+class CloseFailingLocalFileSystem : public arrow::fs::LocalFileSystem {
+ public:
+    using arrow::fs::LocalFileSystem::LocalFileSystem;
+
+    arrow::Result<std::shared_ptr<arrow::io::OutputStream>>
+    OpenOutputStream(const std::string& path,
+                     const std::shared_ptr<const arrow::KeyValueMetadata>&
+                         metadata) override {
+        auto result =
+            arrow::fs::LocalFileSystem::OpenOutputStream(path, metadata);
+        if (!result.ok()) {
+            return result.status();
+        }
+        return std::make_shared<CloseFailingOutputStream>(result.ValueOrDie());
+    }
+};
+
+}  // namespace
+
 class ParquetWriterFactoryTest : public ::testing::Test {
  protected:
     void
@@ -30,6 +107,23 @@ class ParquetWriterFactoryTest : public ::testing::Test {
     std::string path_prefix_;
 };
 
+namespace {
+
+milvus_storage::StorageConfig
+CreatePackedWriterStorageConfig() {
+    return milvus_storage::StorageConfig{};
+}
+
+milvus_storage::ArrowFileSystemPtr
+CreateLocalArrowFileSystem() {
+    milvus::storage::StorageConfig storage_config;
+    storage_config.storage_type = "local";
+    storage_config.root_path = TestLocalPath;
+    return milvus::storage::InitArrowFileSystem(storage_config);
+}
+
+}  // namespace
+
 TEST_F(ParquetWriterFactoryTest, ColumnGroupingStrategyFactoryTest) {
     // Test creating default strategy
     auto default_strategy = ColumnGroupingStrategyFactory::CreateStrategy(
@@ -105,4 +199,65 @@ TEST_F(ParquetWriterFactoryTest, CreateContextWithColumnGroups) {
     EXPECT_EQ(group_ids.size(), column_map_.size());
 }
 
-}  // namespace milvus::index
+TEST_F(ParquetWriterFactoryTest, CloseReturnsStatusAndIsIdempotent) {
+    auto fs = CreateLocalArrowFileSystem();
+    ASSERT_NE(fs, nullptr);
+    auto path_prefix = std::filesystem::path("json_stats_writer_close");
+    auto local_path = std::filesystem::path(TestLocalPath) / path_prefix;
+    std::filesystem::remove_all(local_path);
+    ASSERT_TRUE(std::filesystem::create_directories(local_path));
+
+    std::map<JsonKey, JsonKeyLayoutType> column_map = {
+        {JsonKey("/int", JSONType::INT64), JsonKeyLayoutType::TYPED},
+        {JsonKey("/shared", JSONType::STRING), JsonKeyLayoutType::SHARED},
+    };
+
+    auto storage_config = CreatePackedWriterStorageConfig();
+    JsonStatsParquetWriter writer(fs, storage_config, 16 * 1024 * 1024, 1024);
+    auto context =
+        ParquetWriterFactory::CreateContext(column_map, path_prefix.string());
+    writer.Init(std::move(context));
+
+    writer.AppendValue(JsonKey("/int", JSONType::INT64).ToColumnName(), "42");
+    writer.AppendSharedRow(nullptr, 0);
+    writer.AddCurrentRow();
+
+    auto status = writer.Close();
+    ASSERT_TRUE(status.ok()) << status.ToString();
+    EXPECT_TRUE(writer.Close().ok());
+    EXPECT_FALSE(writer.GetPathsToSize().empty());
+
+    std::filesystem::remove_all(local_path);
+}
+
+TEST_F(ParquetWriterFactoryTest, ClosePropagatesFinalCloseFailure) {
+    auto fs = std::make_shared<CloseFailingLocalFileSystem>();
+    auto path_prefix =
+        std::filesystem::path(TestLocalPath) / "json_stats_writer_close_fail";
+    std::filesystem::remove_all(path_prefix);
+    ASSERT_TRUE(std::filesystem::create_directories(path_prefix));
+
+    std::map<JsonKey, JsonKeyLayoutType> column_map = {
+        {JsonKey("/int", JSONType::INT64), JsonKeyLayoutType::TYPED},
+        {JsonKey("/shared", JSONType::STRING), JsonKeyLayoutType::SHARED},
+    };
+
+    auto storage_config = CreatePackedWriterStorageConfig();
+    JsonStatsParquetWriter writer(fs, storage_config, 16 * 1024 * 1024, 1024);
+    auto context =
+        ParquetWriterFactory::CreateContext(column_map, path_prefix.string());
+    writer.Init(std::move(context));
+
+    writer.AppendValue(JsonKey("/int", JSONType::INT64).ToColumnName(), "42");
+    writer.AppendSharedRow(nullptr, 0);
+    writer.AddCurrentRow();
+
+    auto status = writer.Close();
+    EXPECT_FALSE(status.ok());
+    EXPECT_NE(status.ToString().find("injected final close failure"),
+              std::string::npos);
+
+    std::filesystem::remove_all(path_prefix);
+}
+
+}  // namespace milvus::index