aws: compression: Implement compression of Apache Arrow Parquet

Signed-off-by: Hiroshi Hatake <[email protected]>

Author: cosmo0920

src/aws/compression/arrow/compress.c

@@ -8,6 +8,9 @@
  */
 
 #include <arrow-glib/arrow-glib.h>
+#ifdef FLB_HAVE_ARROW_PARQUET
+#include <parquet-glib/parquet-glib.h>
+#endif
 #include <inttypes.h>
 
 /*
@@ -145,3 +148,120 @@ int out_s3_compress_arrow(void *json, size_t size, void **out_buf, size_t *out_s
         g_bytes_unref(bytes);
         return 0;
 }
+
+#ifdef FLB_HAVE_ARROW_PARQUET
+static GArrowResizableBuffer* table_to_parquet_buffer(GArrowTable *table)
+{
+        GArrowResizableBuffer *buffer;
+        GArrowBufferOutputStream *sink;
+        GParquetArrowFileWriter *writer;
+        GArrowSchema *schema;
+        GError *error = NULL;
+        gboolean success;
+
+        buffer = garrow_resizable_buffer_new(0, &error);
+        if (buffer == NULL) {
+            g_error_free(error);
+            return NULL;
+        }
+
+        sink = garrow_buffer_output_stream_new(buffer);
+        if (sink == NULL) {
+            g_object_unref(buffer);
+            return NULL;
+        }
+
+        schema = garrow_table_get_schema(table);
+        if (schema == NULL) {
+            g_object_unref(buffer);
+            g_object_unref(sink);
+            return NULL;
+        }
+
+        /* Create a new Parquet file writer */
+        writer = gparquet_arrow_file_writer_new_arrow(schema,
+                                                      GARROW_OUTPUT_STREAM(sink),
+                                                      NULL, /* Arrow writer properties */
+                                                      &error);
+        g_object_unref(schema);
+        if (writer == NULL) {
+            g_error_free(error);
+            g_object_unref(buffer);
+            g_object_unref(sink);
+            return NULL;
+        }
+
+        /* Write the entire table to the Parquet file buffer */
+        success = gparquet_arrow_file_writer_write_table(writer, table, 0, &error);
+        if (!success) {
+            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);
+            g_object_unref(buffer);
+            g_object_unref(sink);
+            g_object_unref(writer);
+            return NULL;
+        }
+
+        g_object_unref(sink);
+        g_object_unref(writer);
+        return buffer;
+}
+
+
+int out_s3_compress_parquet(void *json, size_t size, void **out_buf, size_t *out_size)
+{
+        GArrowTable *table;
+        GArrowResizableBuffer *buffer;
+        GBytes *bytes;
+        gconstpointer ptr;
+        gsize len;
+        uint8_t *buf;
+
+        table = parse_json((uint8_t *) json, size);
+        if (table == NULL) {
+            return -1;
+        }
+
+        buffer = table_to_parquet_buffer(table);
+        g_object_unref(table);
+        if (buffer == NULL) {
+            return -1;
+        }
+
+        bytes = garrow_buffer_get_data(GARROW_BUFFER(buffer));
+        if (bytes == NULL) {
+            g_object_unref(buffer);
+            return -1;
+        }
+
+        ptr = g_bytes_get_data(bytes, &len);
+        if (ptr == NULL) {
+            g_object_unref(buffer);
+            g_bytes_unref(bytes);
+            return -1;
+        }
+
+        buf = malloc(len);
+        if (buf == NULL) {
+            g_object_unref(buffer);
+            g_bytes_unref(bytes);
+            return -1;
+        }
+        memcpy(buf, ptr, len);
+        *out_buf = (void *) buf;
+        *out_size = len;
+
+        g_object_unref(buffer);
+        g_bytes_unref(bytes);
+        return 0;
+}
+#endif

src/aws/compression/arrow/compress.h

@@ -11,3 +11,18 @@
  */
 
 int out_s3_compress_arrow(void *json, size_t size, void **out_buf, size_t *out_size);
+
+#ifdef FLB_HAVE_ARROW_PARQUET
+/*
+ * This function converts out_s3 buffer into Apache Parquet format.
+ *
+ * `json` is a string that contain (concatenated) JSON objects.
+ *
+ * `size` is the length of the json data (excluding the trailing
+ * null-terminator character).
+ *
+ * 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);
+#endif