feat(avro): support zstd compression level configuration

Author: dalingmeng

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_);

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_);

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<int> 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));

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<int> compression_level = std::nullopt);
 
     Status AddBatch(ArrowArray* batch) override;
 

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,

src/paimon/format/avro/avro_writer_builder.h

@@ -25,6 +25,7 @@
 #include "avro/DataFile.hh"
 #include "avro/Stream.hh"
 #include "paimon/common/utils/string_utils.h"
+#include "paimon/core/core_options.h"
 #include "paimon/format/avro/avro_format_writer.h"
 #include "paimon/format/avro/avro_output_stream_impl.h"
 #include "paimon/format/writer_builder.h"
@@ -58,7 +59,10 @@ class AvroWriterBuilder : public WriterBuilder {
         auto output_stream = std::make_unique<AvroOutputStreamImpl>(out, BUFFER_SIZE, pool_);
         PAIMON_ASSIGN_OR_RAISE(::avro::Codec codec,
                                ToAvroCompressionKind(StringUtils::ToLowerCase(compression)));
-        return AvroFormatWriter::Create(std::move(output_stream), schema_, codec);
+        PAIMON_ASSIGN_OR_RAISE(std::optional<int> compression_zstd_level,
+                               GetAvroCompressionZstdLevel(codec));
+        return AvroFormatWriter::Create(std::move(output_stream), schema_, codec,
+                                        compression_zstd_level);
     }
 
  private:
@@ -77,6 +81,14 @@ class AvroWriterBuilder : public WriterBuilder {
             return Status::Invalid("unknown compression " + file_compression);
         }
     }
+    Result<std::optional<int>> GetAvroCompressionZstdLevel(const ::avro::Codec& codec) {
+        std::optional<int> compression_zstd_level;
+        if (codec == ::avro::Codec::ZSTD_CODEC) {
+            PAIMON_ASSIGN_OR_RAISE(CoreOptions core_options, CoreOptions::FromMap(options_));
+            compression_zstd_level = core_options.GetFileCompressionZstdLevel();
+        }
+        return compression_zstd_level;
+    }
 
     std::shared_ptr<MemoryPool> pool_;
     std::shared_ptr<arrow::Schema> schema_;

src/paimon/format/avro/avro_writer_builder_test.cpp

@@ -51,4 +51,32 @@ TEST(ToAvroCompressionKindTest, HandlesInvalidCompression) {
 TEST(ToAvroCompressionKindTest, HandlesEmptyString) {
     ASSERT_NOK(AvroWriterBuilder::ToAvroCompressionKind(""));
 }
+
+TEST(ToAvroCompressionKindTest, GetAvroCompressionZstdLevel) {
+    AvroWriterBuilder builder(nullptr, -1, {{Options::FILE_FORMAT, "avro"}});
+    ASSERT_OK_AND_ASSIGN(std::optional<int> zstd_level,
+                         builder.GetAvroCompressionZstdLevel(::avro::Codec::ZSTD_CODEC));
+    ASSERT_TRUE(zstd_level.has_value());
+    ASSERT_EQ(zstd_level.value(), 1);
+
+    ASSERT_OK_AND_ASSIGN(std::optional<int> compression_level1,
+                         builder.GetAvroCompressionZstdLevel(::avro::Codec::SNAPPY_CODEC));
+    ASSERT_FALSE(compression_level1.has_value());
+
+    ASSERT_OK_AND_ASSIGN(std::optional<int> compression_level2,
+                         builder.GetAvroCompressionZstdLevel(::avro::Codec::DEFLATE_CODEC));
+    ASSERT_FALSE(compression_level2.has_value());
+
+    ASSERT_OK_AND_ASSIGN(std::optional<int> compression_level3,
+                         builder.GetAvroCompressionZstdLevel(::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<int> zstd_level2,
+                         builder2.GetAvroCompressionZstdLevel(::avro::Codec::ZSTD_CODEC));
+    ASSERT_TRUE(zstd_level2.has_value());
+    ASSERT_EQ(zstd_level2.value(), 3);
+}
+
 }  // namespace paimon::avro::test

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