diff --git a/src/aws/compression/arrow/compress.c b/src/aws/compression/arrow/compress.c index 16eb39d3df6..0ebf23a72e9 100644 --- a/src/aws/compression/arrow/compress.c +++ b/src/aws/compression/arrow/compress.c @@ -39,6 +39,34 @@ static GArrowCompressionType compression_type_to_garrow(int compression_type) } } +static int validate_columnar_compression(int columnar_format, + int compression_type) +{ + if (columnar_format == FLB_AWS_COMPRESS_FORMAT_PARQUET) { + switch (compression_type) { + case FLB_AWS_COMPRESS_NONE: + case FLB_AWS_COMPRESS_SNAPPY: + case FLB_AWS_COMPRESS_GZIP: + case FLB_AWS_COMPRESS_ZSTD: + return 0; + default: + return -1; + } + } + + if (columnar_format == FLB_AWS_COMPRESS_FORMAT_ARROW) { + switch (compression_type) { + case FLB_AWS_COMPRESS_NONE: + case FLB_AWS_COMPRESS_ZSTD: + return 0; + default: + return -1; + } + } + + return -1; +} + static int choose_block_size(size_t size) { int block_size = 8 * 1024 * 1024; @@ -251,6 +279,14 @@ int flb_aws_compression_compress_columnar(int columnar_format, gsize len; uint8_t *buf; + if (validate_columnar_compression(columnar_format, + compression_type) != 0) { + flb_error("[aws][compress] unsupported compression type %d " + "for columnar format %d", + compression_type, columnar_format); + return -1; + } + table = parse_json((uint8_t *) json, size); if (table == NULL) { flb_error("[aws][compress] Failed to parse JSON into Arrow Table"); diff --git a/tests/internal/aws_compress.c b/tests/internal/aws_compress.c index e081a7f433a..49b67b0ffb7 100644 --- a/tests/internal/aws_compress.c +++ b/tests/internal/aws_compress.c @@ -419,6 +419,22 @@ void test_parquet_format_uncompressed() flb_free(out_buf); } +void test_parquet_format_invalid_compression() +{ + int ret; + void *out_buf = NULL; + size_t out_size = 0; + char *json = "{\"key\":\"value\",\"num\":42}\n"; + size_t json_len = strlen(json); + + ret = flb_aws_compression_compress_columnar(FLB_AWS_COMPRESS_FORMAT_PARQUET, + json, json_len, + &out_buf, &out_size, -1); + TEST_CHECK(ret == -1); + TEST_CHECK(out_buf == NULL); + TEST_CHECK(out_size == 0); +} + void test_parquet_compression_reduces_size() { int ret; @@ -545,6 +561,8 @@ TEST_LIST = { { "test_parquet_format_zstd", test_parquet_format_zstd }, { "test_parquet_format_gzip", test_parquet_format_gzip }, { "test_parquet_format_uncompressed", test_parquet_format_uncompressed }, + { "test_parquet_format_invalid_compression", + test_parquet_format_invalid_compression }, { "test_parquet_compression_reduces_size", test_parquet_compression_reduces_size }, #endif