aws: add page-level compression support to parquet writer

The Parquet writer passed NULL for GParquetWriterProperties, causing
all output files to have UNCOMPRESSED pages. This adds a
parquet_compression parameter that configures the page-level codec
(snappy, gzip, zstd, or none) via the GLib Parquet writer properties
API.

The parquet entry is removed from the compression dispatch table since
the new function signature includes the codec parameter and no longer
matches the generic compress function pointer type. The S3 plugin now
calls out_s3_compress_parquet directly.

Signed-off-by: Rituparna Khaund <ritukhau@amazon.co.uk>

Author: Rituparna Khaund

include/fluent-bit/aws/flb_aws_compress.h

@@ -28,6 +28,12 @@
 #define FLB_AWS_COMPRESS_ZSTD    4
 #define FLB_AWS_COMPRESS_SNAPPY  5
 
+/* Parquet page-level compression codecs */
+#define FLB_PARQUET_COMPRESSION_NONE    0
+#define FLB_PARQUET_COMPRESSION_SNAPPY  1
+#define FLB_PARQUET_COMPRESSION_GZIP    2
+#define FLB_PARQUET_COMPRESSION_ZSTD    3
+
 /*
  * Get compression type from compression keyword. The return value is used to identify
  * what compression option to utilize.
@@ -63,4 +69,13 @@ int flb_aws_compression_b64_truncate_compress(int compression_type, size_t max_o
                                              void *in_data, size_t in_len,
                                              void **out_data, size_t *out_len);
 
+/*
+ * Convert JSON data to Apache Parquet format with the specified
+ * page-level compression codec (FLB_PARQUET_COMPRESSION_*).
+ * Returns 0 on success, -1 on failure.
+ */
+int out_s3_compress_parquet(void *json, size_t size,
+                            void **out_buf, size_t *out_size,
+                            int parquet_compression);
+
 #endif

src/aws/compression/arrow/compress.c

@@ -170,11 +170,29 @@ int out_s3_compress_arrow(void *json, size_t size, void **out_buf, size_t *out_s
 }
 
 #ifdef FLB_HAVE_ARROW_PARQUET
-static GArrowResizableBuffer* table_to_parquet_buffer(GArrowTable *table)
+#include <fluent-bit/aws/flb_aws_compress.h>
+
+static GArrowCompressionType parquet_compression_to_garrow(int parquet_compression)
+{
+        switch (parquet_compression) {
+        case FLB_PARQUET_COMPRESSION_SNAPPY:
+            return GARROW_COMPRESSION_TYPE_SNAPPY;
+        case FLB_PARQUET_COMPRESSION_GZIP:
+            return GARROW_COMPRESSION_TYPE_GZIP;
+        case FLB_PARQUET_COMPRESSION_ZSTD:
+            return GARROW_COMPRESSION_TYPE_ZSTD;
+        default:
+            return GARROW_COMPRESSION_TYPE_UNCOMPRESSED;
+        }
+}
+
+static GArrowResizableBuffer* table_to_parquet_buffer(GArrowTable *table,
+                                                      int parquet_compression)
 {
         GArrowResizableBuffer *buffer;
         GArrowBufferOutputStream *sink;
         GParquetArrowFileWriter *writer;
+        GParquetWriterProperties *props;
         GArrowSchema *schema;
         GError *error = NULL;
         gboolean success;
@@ -199,14 +217,19 @@ static GArrowResizableBuffer* table_to_parquet_buffer(GArrowTable *table)
             return NULL;
         }
 
-        /* Create a new Parquet file writer */
+        props = gparquet_writer_properties_new();
+        gparquet_writer_properties_set_compression(
+            props, parquet_compression_to_garrow(parquet_compression), NULL);
+
         writer = gparquet_arrow_file_writer_new_arrow(schema,
                                                       GARROW_OUTPUT_STREAM(sink),
-                                                      NULL, /* Arrow writer properties */
+                                                      props,
                                                       &error);
         g_object_unref(schema);
+        g_object_unref(props);
         if (writer == NULL) {
-            flb_error("[aws][compress] Failed to create parquet writer: %s", error->message);
+            flb_error("[aws][compress] Failed to create parquet writer: %s",
+                      error->message);
             g_error_free(error);
             g_object_unref(buffer);
             g_object_unref(sink);
@@ -215,18 +238,18 @@ static GArrowResizableBuffer* table_to_parquet_buffer(GArrowTable *table)
 
         n_rows = garrow_table_get_n_rows(table);
 
-        /* Write the entire table to the Parquet file buffer */
-        success = gparquet_arrow_file_writer_write_table(writer, table, n_rows, &error);
+        success = gparquet_arrow_file_writer_write_table(writer, table,
+                                                         n_rows, &error);
         if (!success) {
-            flb_error("[aws][compress] Failed to write table to parquet buffer: %s", error->message);
+            flb_error("[aws][compress] Failed to write table to parquet "
+                      "buffer: %s", error->message);
             g_error_free(error);
             g_object_unref(buffer);
             g_object_unref(sink);
             g_object_unref(writer);
             return NULL;
         }
 
-        /* Close the writer to finalize the Parquet file metadata */
         success = gparquet_arrow_file_writer_close(writer, &error);
         if (!success) {
             g_error_free(error);
@@ -242,7 +265,8 @@ static GArrowResizableBuffer* table_to_parquet_buffer(GArrowTable *table)
 }
 
 
-int out_s3_compress_parquet(void *json, size_t size, void **out_buf, size_t *out_size)
+int out_s3_compress_parquet(void *json, size_t size, void **out_buf,
+                           size_t *out_size, int parquet_compression)
 {
         GArrowTable *table;
         GArrowResizableBuffer *buffer;
@@ -253,14 +277,16 @@ int out_s3_compress_parquet(void *json, size_t size, void **out_buf, size_t *out
 
         table = parse_json((uint8_t *) json, size);
         if (table == NULL) {
-            flb_error("[aws][compress] Failed to parse JSON into Arrow Table for Parquet conversion");
+            flb_error("[aws][compress] Failed to parse JSON into Arrow Table"
+                      " for Parquet conversion");
             return -1;
         }
 
-        buffer = table_to_parquet_buffer(table);
+        buffer = table_to_parquet_buffer(table, parquet_compression);
         g_object_unref(table);
         if (buffer == NULL) {
-            flb_error("[aws][compress] Failed to convert Arrow Table into Parquet buffer");
+            flb_error("[aws][compress] Failed to convert Arrow Table into"
+                      " Parquet buffer");
             return -1;
         }
 

src/aws/compression/arrow/compress.h

@@ -21,8 +21,12 @@ int out_s3_compress_arrow(void *json, size_t size, void **out_buf, size_t *out_s
  * `size` is the length of the json data (excluding the trailing
  * null-terminator character).
  *
+ * `parquet_compression` specifies the page-level codec to use
+ * (FLB_PARQUET_COMPRESSION_*).
+ *
  * Return 0 on success (with `out_buf` and `out_size` updated),
  * and -1 on failure
  */
-int out_s3_compress_parquet(void *json, size_t size, void **out_buf, size_t *out_size);
+int out_s3_compress_parquet(void *json, size_t size, void **out_buf,
+                            size_t *out_size, int parquet_compression);
 #endif

src/aws/flb_aws_compress.c

@@ -83,13 +83,6 @@ static const struct compression_option compression_options[] = {
         "arrow",
         &out_s3_compress_arrow
     },
-#endif
-#ifdef FLB_HAVE_ARROW_PARQUET
-    {
-        FLB_AWS_COMPRESS_PARQUET,
-        "parquet",
-        &out_s3_compress_parquet
-    },
 #endif
     { 0 }
 };