Persist aggregate state settings for plain exports

Carry the Parquet and Iceberg aggregate-state gates through plain `MergeTree` export task serialization and worker contexts. Keep older descriptors compatible by treating missing fields as disabled.

Author: zvonand

src/Storages/MergeTree/MergeTreePartitionExportTask.h

@@ -82,6 +82,8 @@ struct MergeTreePartitionExportTask
     String filename_pattern;
     bool write_full_path_in_iceberg_metadata = false;
     bool allow_lossy_cast = false;
+    bool allow_aggregate_function_states_in_parquet = false;
+    bool allow_aggregate_function_states_in_iceberg = false;
     String iceberg_metadata_json;
 
     /// Optional for backwards compatibility with descriptors written before these
@@ -189,6 +191,8 @@ struct MergeTreePartitionExportTask
         json.set("filename_pattern", filename_pattern);
         json.set("write_full_path_in_iceberg_metadata", write_full_path_in_iceberg_metadata);
         json.set("allow_lossy_cast", allow_lossy_cast);
+        json.set("allow_aggregate_function_states_in_parquet", allow_aggregate_function_states_in_parquet);
+        json.set("allow_aggregate_function_states_in_iceberg", allow_aggregate_function_states_in_iceberg);
         if (!iceberg_metadata_json.empty())
             json.set("iceberg_metadata_json", iceberg_metadata_json);
         if (parquet_compression_method)
@@ -281,6 +285,10 @@ struct MergeTreePartitionExportTask
         task.filename_pattern = json->getValue<String>("filename_pattern");
         task.write_full_path_in_iceberg_metadata = json->getValue<bool>("write_full_path_in_iceberg_metadata");
         task.allow_lossy_cast = json->getValue<bool>("allow_lossy_cast");
+        if (json->has("allow_aggregate_function_states_in_parquet"))
+            task.allow_aggregate_function_states_in_parquet = json->getValue<bool>("allow_aggregate_function_states_in_parquet");
+        if (json->has("allow_aggregate_function_states_in_iceberg"))
+            task.allow_aggregate_function_states_in_iceberg = json->getValue<bool>("allow_aggregate_function_states_in_iceberg");
         if (json->has("iceberg_metadata_json"))
             task.iceberg_metadata_json = json->getValue<String>("iceberg_metadata_json");
         if (json->has("parquet_compression_method"))

src/Storages/MergeTree/tests/gtest_export_partition_ordering.cpp

@@ -3,6 +3,7 @@
 #include <sstream>
 #include <Storages/ExportReplicatedMergeTreePartitionTaskEntry.h>
 #include <Storages/MergeTree/ExportPartitionUtils.h>
+#include <Storages/MergeTree/MergeTreePartitionExportTask.h>
 #include <Common/tests/gtest_global_context.h>
 #include <Core/Settings.h>
 #include <Poco/JSON/Object.h>
@@ -51,6 +52,29 @@ namespace
         manifest.filename_pattern = "{part_name}";
         return manifest;
     }
+
+    MergeTreePartitionExportTask makeValidPlainTask()
+    {
+        MergeTreePartitionExportTask task;
+        task.transaction_id = "tx1";
+        task.query_id = "query1";
+        task.partition_id = "2020";
+        task.source_database = "db1";
+        task.source_table = "source1";
+        task.destination_database = "db1";
+        task.destination_table = "table1";
+        task.create_time = 1000;
+        task.parts.push_back({"part1", false, {}});
+        task.task_timeout_seconds = 60;
+        task.max_threads = 1;
+        task.parallel_formatting = true;
+        task.parquet_parallel_encoding = true;
+        task.max_bytes_per_file = 1000000;
+        task.max_rows_per_file = 1000;
+        task.file_already_exists_policy = MergeTreePartExportManifest::FileAlreadyExistsPolicy::error;
+        task.filename_pattern = "{part_name}";
+        return task;
+    }
 }
 
 class ExportPartitionOrderingTest : public ::testing::Test
@@ -270,6 +294,47 @@ TEST_F(ExportPartitionManifestBackCompatTest, AggregateFunctionStateGatesApplied
     }
 }
 
+TEST_F(ExportPartitionManifestBackCompatTest, PlainTaskMissingAggregateFunctionStateGatesParseAsDisabled)
+{
+    auto task = makeValidPlainTask();
+    task.allow_aggregate_function_states_in_parquet = true;
+    task.allow_aggregate_function_states_in_iceberg = true;
+
+    Poco::JSON::Parser parser;
+    auto json = parser.parse(task.toJsonString()).extract<Poco::JSON::Object::Ptr>();
+    json->remove("allow_aggregate_function_states_in_parquet");
+    json->remove("allow_aggregate_function_states_in_iceberg");
+    std::ostringstream oss;     // STYLE_CHECK_ALLOW_STD_STRING_STREAM
+    oss.exceptions(std::ios::failbit);
+    Poco::JSON::Stringifier::stringify(json, oss);
+
+    auto parsed = MergeTreePartitionExportTask::fromJsonString(oss.str());
+    EXPECT_FALSE(parsed.allow_aggregate_function_states_in_parquet);
+    EXPECT_FALSE(parsed.allow_aggregate_function_states_in_iceberg);
+}
+
+TEST_F(ExportPartitionManifestBackCompatTest, PlainTaskAggregateFunctionStateGatesAppliedToWorkerContextForEveryValue)
+{
+    for (const bool value : {false, true})
+    {
+        auto task = makeValidPlainTask();
+        task.allow_aggregate_function_states_in_parquet = value;
+        task.allow_aggregate_function_states_in_iceberg = value;
+
+        auto parsed = MergeTreePartitionExportTask::fromJsonString(task.toJsonString());
+        auto worker_context = ExportPartitionUtils::getContextCopyWithTaskSettings(getContext().context, parsed);
+
+        EXPECT_EQ(parsed.allow_aggregate_function_states_in_parquet, value) << "value=" << value;
+        EXPECT_EQ(parsed.allow_aggregate_function_states_in_iceberg, value) << "value=" << value;
+        EXPECT_EQ(
+            worker_context->getSettingsRef()[Setting::allow_experimental_aggregate_function_states_in_parquet].value,
+            value) << "value=" << value;
+        EXPECT_EQ(
+            worker_context->getSettingsRef()[Setting::allow_experimental_aggregate_function_states_in_iceberg].value,
+            value) << "value=" << value;
+    }
+}
+
 TEST(ExportPartitionRetryClassification, MissingPartIsFatalOnlyOnPlainPath)
 {
     EXPECT_FALSE(ExportPartitionUtils::isNonRetryableExportError(ErrorCodes::NO_SUCH_DATA_PART));

src/Storages/StorageMergeTree.cpp

@@ -138,6 +138,8 @@ namespace Setting
     extern const SettingsTimezone iceberg_partition_timezone;
     extern const SettingsMergeTreePartExportSchemaMatchMode export_merge_tree_part_schema_match_mode;
     extern const SettingsBool export_merge_tree_part_ignore_extra_source_columns;
+    extern const SettingsBool allow_experimental_aggregate_function_states_in_parquet;
+    extern const SettingsBool allow_experimental_aggregate_function_states_in_iceberg;
 }
 
 namespace ServerSetting
@@ -3894,6 +3896,10 @@ void StorageMergeTree::exportPartitionToTable(const PartitionCommand & command,
     descriptor.filename_pattern = query_context->getSettingsRef()[Setting::export_merge_tree_part_filename_pattern].value;
     descriptor.write_full_path_in_iceberg_metadata = query_context->getSettingsRef()[Setting::write_full_path_in_iceberg_metadata];
     descriptor.allow_lossy_cast = query_context->getSettingsRef()[Setting::export_merge_tree_part_allow_lossy_cast];
+    descriptor.allow_aggregate_function_states_in_parquet
+        = query_context->getSettingsRef()[Setting::allow_experimental_aggregate_function_states_in_parquet];
+    descriptor.allow_aggregate_function_states_in_iceberg
+        = query_context->getSettingsRef()[Setting::allow_experimental_aggregate_function_states_in_iceberg];
     descriptor.iceberg_partition_timezone = query_context->getSettingsRef()[Setting::iceberg_partition_timezone].toString();
     descriptor.schema_match_mode = query_context->getSettingsRef()[Setting::export_merge_tree_part_schema_match_mode].value;
     descriptor.ignore_extra_source_columns = query_context->getSettingsRef()[Setting::export_merge_tree_part_ignore_extra_source_columns].value;