feat(avro): support zstd compression level configuration

dalingmeng · dalingmeng · commit c67750dced64 · 2026-02-11T17:48:36.000+08:00
diff --git a/src/paimon/format/avro/avro_file_batch_reader_test.cpp b/src/paimon/format/avro/avro_file_batch_reader_test.cpp
@@ -39,7 +39,8 @@ namespace paimon::avro::test {
 class AvroFileBatchReaderTest : public ::testing::Test, public ::testing::WithParamInterface<bool> {
  public:
     void SetUp() override {
-        ASSERT_OK_AND_ASSIGN(file_format_, FileFormatFactory::Get("avro", {}));
+        ASSERT_OK_AND_ASSIGN(file_format_,
+                             FileFormatFactory::Get("avro", {{Options::FILE_FORMAT, "avro"}}));
         fs_ = std::make_shared<LocalFileSystem>();
         dir_ = ::paimon::test::UniqueTestDirectory::Create();
         ASSERT_TRUE(dir_);
diff --git a/src/paimon/format/avro/avro_file_format_test.cpp b/src/paimon/format/avro/avro_file_format_test.cpp
@@ -45,7 +45,8 @@ namespace paimon::avro::test {
 class AvroFileFormatTest : public testing::Test, public ::testing::WithParamInterface<std::string> {
  public:
     void SetUp() override {
-        ASSERT_OK_AND_ASSIGN(file_format_, FileFormatFactory::Get("avro", {}));
+        ASSERT_OK_AND_ASSIGN(file_format_,
+                             FileFormatFactory::Get("avro", {{Options::FILE_FORMAT, "avro"}}));
         fs_ = std::make_shared<LocalFileSystem>();
         dir_ = ::paimon::test::UniqueTestDirectory::Create();
         ASSERT_TRUE(dir_);
diff --git a/src/paimon/format/avro/avro_format_defs.h b/src/paimon/format/avro/avro_format_defs.h
@@ -0,0 +1,27 @@
+/*
+ * Copyright 2026-present Alibaba Inc.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+#pragma once
+
+#include <cstdint>
+#include <limits>
+namespace paimon::avro {
+
+// write
+static constexpr int32_t AVRO_BUFFER_SIZE = 1024 * 1024;
+static inline const char AVRO_CODEC[] = "avro.codec";
+
+}  // namespace paimon::avro
diff --git a/src/paimon/format/avro/avro_format_writer.cpp b/src/paimon/format/avro/avro_format_writer.cpp
@@ -53,13 +53,14 @@ AvroFormatWriter::AvroFormatWriter(std::unique_ptr<::avro::DataFileWriterBase>&&
 
 Result<std::unique_ptr<AvroFormatWriter>> AvroFormatWriter::Create(
     std::unique_ptr<AvroOutputStreamImpl> out, const std::shared_ptr<arrow::Schema>& schema,
-    const ::avro::Codec codec) {
+    const ::avro::Codec codec, std::optional<int32_t> compression_level) {
     try {
         PAIMON_ASSIGN_OR_RAISE(::avro::ValidSchema avro_schema,
                                AvroSchemaConverter::ArrowSchemaToAvroSchema(schema));
         AvroOutputStreamImpl* avro_output_stream = out.get();
-        auto writer = std::make_unique<::avro::DataFileWriterBase>(std::move(out), avro_schema,
-                                                                   DEFAULT_SYNC_INTERVAL, codec);
+        auto writer = std::make_unique<::avro::DataFileWriterBase>(
+            std::move(out), avro_schema, DEFAULT_SYNC_INTERVAL, codec, ::avro::Metadata(),
+            compression_level);
         auto data_type = arrow::struct_(schema->fields());
         return std::unique_ptr<AvroFormatWriter>(
             new AvroFormatWriter(std::move(writer), avro_schema, data_type, avro_output_stream));
diff --git a/src/paimon/format/avro/avro_format_writer.h b/src/paimon/format/avro/avro_format_writer.h
@@ -19,6 +19,7 @@
 #include <cstddef>
 #include <cstdint>
 #include <memory>
+#include <optional>
 
 #include "arrow/api.h"
 #include "avro/DataFile.hh"
@@ -49,7 +50,7 @@ class AvroFormatWriter : public FormatWriter {
  public:
     static Result<std::unique_ptr<AvroFormatWriter>> Create(
         std::unique_ptr<AvroOutputStreamImpl> out, const std::shared_ptr<arrow::Schema>& schema,
-        const ::avro::Codec codec);
+        const ::avro::Codec codec, std::optional<int32_t> compression_level = std::nullopt);
 
     Status AddBatch(ArrowArray* batch) override;
 
diff --git a/src/paimon/format/avro/avro_format_writer_test.cpp b/src/paimon/format/avro/avro_format_writer_test.cpp
@@ -68,7 +68,8 @@ class AvroFormatWriterTest : public ::testing::Test {
                                                      int32_t batch_size) {
         ::ArrowSchema c_schema;
         EXPECT_TRUE(arrow::ExportSchema(*schema, &c_schema).ok());
-        EXPECT_OK_AND_ASSIGN(auto file_format, FileFormatFactory::Get("avro", {}));
+        EXPECT_OK_AND_ASSIGN(auto file_format,
+                             FileFormatFactory::Get("avro", {{Options::FILE_FORMAT, "avro"}}));
         EXPECT_OK_AND_ASSIGN(auto writer_builder,
                              file_format->CreateWriterBuilder(&c_schema, batch_size));
         EXPECT_OK_AND_ASSIGN(std::shared_ptr<FormatWriter> writer,
diff --git a/src/paimon/format/avro/avro_writer_builder.h b/src/paimon/format/avro/avro_writer_builder.h
@@ -24,7 +24,10 @@
 
 #include "avro/DataFile.hh"
 #include "avro/Stream.hh"
+#include "paimon/common/utils/options_utils.h"
 #include "paimon/common/utils/string_utils.h"
+#include "paimon/core/core_options.h"
+#include "paimon/format/avro/avro_format_defs.h"
 #include "paimon/format/avro/avro_format_writer.h"
 #include "paimon/format/avro/avro_output_stream_impl.h"
 #include "paimon/format/writer_builder.h"
@@ -55,15 +58,19 @@ class AvroWriterBuilder : public WriterBuilder {
 
     Result<std::unique_ptr<FormatWriter>> Build(const std::shared_ptr<OutputStream>& out,
                                                 const std::string& compression) override {
-        auto output_stream = std::make_unique<AvroOutputStreamImpl>(out, BUFFER_SIZE, pool_);
+        auto output_stream = std::make_unique<AvroOutputStreamImpl>(out, AVRO_BUFFER_SIZE, pool_);
+        PAIMON_ASSIGN_OR_RAISE(
+            std::string file_compression,
+            OptionsUtils::GetValueFromMap<std::string>(options_, AVRO_CODEC, compression));
         PAIMON_ASSIGN_OR_RAISE(::avro::Codec codec,
-                               ToAvroCompressionKind(StringUtils::ToLowerCase(compression)));
-        return AvroFormatWriter::Create(std::move(output_stream), schema_, codec);
+                               ToAvroCompressionKind(StringUtils::ToLowerCase(file_compression)));
+        PAIMON_ASSIGN_OR_RAISE(std::optional<int32_t> compression_level,
+                               GetAvroCompressionLevel(codec));
+        return AvroFormatWriter::Create(std::move(output_stream), schema_, codec,
+                                        compression_level);
     }
 
  private:
-    static constexpr int32_t BUFFER_SIZE = 1024 * 1024;
-
     static Result<::avro::Codec> ToAvroCompressionKind(const std::string& file_compression) {
         if (file_compression == "zstd" || file_compression == "zstandard") {
             return ::avro::Codec::ZSTD_CODEC;
@@ -77,6 +84,14 @@ class AvroWriterBuilder : public WriterBuilder {
             return Status::Invalid("unknown compression " + file_compression);
         }
     }
+    Result<std::optional<int32_t>> GetAvroCompressionLevel(const ::avro::Codec& codec) {
+        std::optional<int32_t> compression_level;
+        if (codec == ::avro::Codec::ZSTD_CODEC) {
+            PAIMON_ASSIGN_OR_RAISE(CoreOptions core_options, CoreOptions::FromMap(options_));
+            compression_level = core_options.GetFileCompressionZstdLevel();
+        }
+        return compression_level;
+    }
 
     std::shared_ptr<MemoryPool> pool_;
     std::shared_ptr<arrow::Schema> schema_;
diff --git a/src/paimon/format/avro/avro_writer_builder_test.cpp b/src/paimon/format/avro/avro_writer_builder_test.cpp
@@ -22,7 +22,7 @@
 
 namespace paimon::avro::test {
 
-TEST(ToAvroCompressionKindTest, HandlesValidCompressions) {
+TEST(AvroWriterBuilderTest, HandlesValidCompressions) {
     ASSERT_OK_AND_ASSIGN(::avro::Codec zstd_codec,
                          AvroWriterBuilder::ToAvroCompressionKind("zstd"));
     ASSERT_EQ(zstd_codec, ::avro::Codec::ZSTD_CODEC);
@@ -44,11 +44,99 @@ TEST(ToAvroCompressionKindTest, HandlesValidCompressions) {
     ASSERT_EQ(deflate_codec, ::avro::Codec::DEFLATE_CODEC);
 }
 
-TEST(ToAvroCompressionKindTest, HandlesInvalidCompression) {
+TEST(AvroWriterBuilderTest, HandlesInvalidCompression) {
     ASSERT_NOK(AvroWriterBuilder::ToAvroCompressionKind("unknown_compression"));
 }
 
-TEST(ToAvroCompressionKindTest, HandlesEmptyString) {
+TEST(AvroWriterBuilderTest, HandlesEmptyString) {
     ASSERT_NOK(AvroWriterBuilder::ToAvroCompressionKind(""));
 }
+
+TEST(AvroWriterBuilderTest, CheckAvroCodec) {
+    arrow::FieldVector fields = {arrow::field("f0", arrow::int32())};
+    auto schema = std::make_shared<arrow::Schema>(fields);
+    {
+        AvroWriterBuilder builder(schema, -1,
+                                  {{Options::FILE_FORMAT, "avro"}, {"avro.codec", "snappy"}});
+        ASSERT_OK_AND_ASSIGN(auto file_writer, builder.Build(nullptr, "zstd"));
+        auto* avro_file_writer = dynamic_cast<AvroFormatWriter*>(file_writer.get());
+        ASSERT_EQ(avro_file_writer->writer_->codec_, ::avro::Codec::SNAPPY_CODEC);
+        ASSERT_EQ(avro_file_writer->writer_->compressionLevel_, std::nullopt);
+    }
+    {
+        AvroWriterBuilder builder2(schema, -1,
+                                   {{Options::FILE_FORMAT, "avro"}, {"avro.codec", "deflate"}});
+        ASSERT_OK_AND_ASSIGN(auto file_writer2, builder2.Build(nullptr, "zstd"));
+        auto* avro_file_writer2 = dynamic_cast<AvroFormatWriter*>(file_writer2.get());
+        ASSERT_EQ(avro_file_writer2->writer_->codec_, ::avro::Codec::DEFLATE_CODEC);
+        ASSERT_EQ(avro_file_writer2->writer_->compressionLevel_, std::nullopt);
+    }
+    {
+        AvroWriterBuilder builder3(schema, -1,
+                                   {{Options::FILE_FORMAT, "avro"}, {"avro.codec", "zstd"}});
+        ASSERT_OK_AND_ASSIGN(auto file_writer3, builder3.Build(nullptr, "zstd"));
+        auto* avro_file_writer3 = dynamic_cast<AvroFormatWriter*>(file_writer3.get());
+        ASSERT_EQ(avro_file_writer3->writer_->codec_, ::avro::Codec::ZSTD_CODEC);
+        ASSERT_EQ(avro_file_writer3->writer_->compressionLevel_, 1);
+    }
+    {
+        AvroWriterBuilder builder4(schema, -1,
+                                   {{Options::FILE_FORMAT, "avro"},
+                                    {"avro.codec", "zstd"},
+                                    {Options::FILE_COMPRESSION_ZSTD_LEVEL, "3"}});
+        ASSERT_OK_AND_ASSIGN(auto file_writer4, builder4.Build(nullptr, "zstd"));
+        auto* avro_file_writer4 = dynamic_cast<AvroFormatWriter*>(file_writer4.get());
+        ASSERT_EQ(avro_file_writer4->writer_->codec_, ::avro::Codec::ZSTD_CODEC);
+        ASSERT_EQ(avro_file_writer4->writer_->compressionLevel_, 3);
+    }
+    {
+        AvroWriterBuilder builder5(schema, -1,
+                                   {{Options::FILE_FORMAT, "avro"},
+                                    {"avro.codec", "null"},
+                                    {Options::FILE_COMPRESSION_ZSTD_LEVEL, "3"}});
+        ASSERT_OK_AND_ASSIGN(auto file_writer5, builder5.Build(nullptr, "zstd"));
+        auto* avro_file_writer5 = dynamic_cast<AvroFormatWriter*>(file_writer5.get());
+        ASSERT_EQ(avro_file_writer5->writer_->codec_, ::avro::Codec::NULL_CODEC);
+        ASSERT_EQ(avro_file_writer5->writer_->compressionLevel_, std::nullopt);
+    }
+    {
+        AvroWriterBuilder builder6(schema, -1,
+                                   {{Options::FILE_FORMAT, "avro"},
+                                    {"avro.codec", "test"},
+                                    {Options::FILE_COMPRESSION_ZSTD_LEVEL, "3"}});
+        ASSERT_NOK(builder6.Build(nullptr, "zstd"));
+    }
+}
+
+TEST(AvroWriterBuilderTest, CheckAvroCompressionLevel) {
+    {
+        AvroWriterBuilder builder(nullptr, -1, {{Options::FILE_FORMAT, "avro"}});
+        ASSERT_OK_AND_ASSIGN(std::optional<int32_t> zstd_level,
+                             builder.GetAvroCompressionLevel(::avro::Codec::ZSTD_CODEC));
+        ASSERT_TRUE(zstd_level.has_value());
+        ASSERT_EQ(zstd_level.value(), 1);
+
+        ASSERT_OK_AND_ASSIGN(std::optional<int32_t> compression_level1,
+                             builder.GetAvroCompressionLevel(::avro::Codec::SNAPPY_CODEC));
+        ASSERT_FALSE(compression_level1.has_value());
+
+        ASSERT_OK_AND_ASSIGN(std::optional<int32_t> compression_level2,
+                             builder.GetAvroCompressionLevel(::avro::Codec::DEFLATE_CODEC));
+        ASSERT_FALSE(compression_level2.has_value());
+
+        ASSERT_OK_AND_ASSIGN(std::optional<int32_t> compression_level3,
+                             builder.GetAvroCompressionLevel(::avro::Codec::NULL_CODEC));
+        ASSERT_FALSE(compression_level3.has_value());
+    }
+    {
+        AvroWriterBuilder builder2(
+            nullptr, -1,
+            {{Options::FILE_FORMAT, "avro"}, {Options::FILE_COMPRESSION_ZSTD_LEVEL, "3"}});
+        ASSERT_OK_AND_ASSIGN(std::optional<int32_t> zstd_level2,
+                             builder2.GetAvroCompressionLevel(::avro::Codec::ZSTD_CODEC));
+        ASSERT_TRUE(zstd_level2.has_value());
+        ASSERT_EQ(zstd_level2.value(), 3);
+    }
+}
+
 }  // namespace paimon::avro::test
diff --git a/third_party/versions.txt b/third_party/versions.txt
@@ -55,8 +55,8 @@ PAIMON_GTEST_PKG_NAME=gtest-${PAIMON_GTEST_BUILD_VERSION}.tar.gz
 PAIMON_ARROW_BUILD_VERSION=17.0.0
 PAIMON_ARROW_BUILD_SHA256_CHECKSUM=9d280d8042e7cf526f8c28d170d93bfab65e50f94569f6a790982a878d8d898d
 PAIMON_ARROW_PKG_NAME=apache-arrow-${PAIMON_ARROW_BUILD_VERSION}.tar.gz
-PAIMON_AVRO_BUILD_VERSION=54b332161524086dcb6cde8afe097097eed7f3ee
-PAIMON_AVRO_BUILD_SHA256_CHECKSUM=00febd590b1e328d3a97b67a6d29a1d0243e0e41bb2b1582ec580d37698d1fe2
+PAIMON_AVRO_BUILD_VERSION=c499eefb48aa2db906c7bca14a047223806f36db
+PAIMON_AVRO_BUILD_SHA256_CHECKSUM=9771f1dcfe3c01aff7ff670e873e66d3406362f71941821d482de65f3d32d780
 PAIMON_AVRO_PKG_NAME=avro-${PAIMON_AVRO_BUILD_VERSION}.tar.gz
 PAIMON_FMT_BUILD_VERSION=11.2.0
 PAIMON_FMT_BUILD_SHA256_CHECKSUM=bc23066d87ab3168f27cef3e97d545fa63314f5c79df5ea444d41d56f962c6af
