Merge pull request #185 from protegrity/av_sequncer_decryption_023

Implement CompressionAndJoin to complete the Decryption sequence

Author: avalerio-tkd

src/server/encryption_sequencer.cpp

@@ -208,15 +208,13 @@ bool DataBatchEncryptionSequencer::ConvertAndDecrypt(const std::vector<uint8_t>&
         auto [encrypted_level_bytes, encrypted_value_bytes] = SplitWithLengthPrefix(decompressed_encrypted_bytes);
         auto level_bytes = encryptor_->DecryptBlock(encrypted_level_bytes);
         auto typed_list = encryptor_->DecryptValueList(encrypted_value_bytes);
-
-
-        throw DBPSUnsupportedException("Per-value decryption implementation is not complete.");
 
         // Convert the decrypted typed list back to value bytes
         auto value_bytes = GetTypedListAsValueBytes(typed_list, datatype_, datatype_length_, format_);
 
         // Join the decrypted level and value bytes, then compress to get plaintext
-        decrypted_result_ = CompressAndJoin(level_bytes, value_bytes);
+        decrypted_result_ = CompressAndJoin(
+            level_bytes, value_bytes, compression_, encoding_attributes_converted_);
     }
 
     // Per-block encryption

src/server/parquet_utils.cpp

@@ -277,8 +277,41 @@ LevelAndValueBytes DecompressAndSplit(
 
 std::vector<uint8_t> CompressAndJoin(
     const std::vector<uint8_t>& level_bytes,
-    const std::vector<uint8_t>& value_bytes) {
-    throw DBPSUnsupportedException("CompressAndJoin not implemented");
+    const std::vector<uint8_t>& value_bytes,
+    CompressionCodec::type compression,
+    const AttributesMap& encoding_attributes) {
+
+    // Get the page type from the encoding attributes.
+    const auto& page_type = std::get<std::string>(encoding_attributes.at("page_type"));
+    
+    // Check that the calculated level bytes size == the size of the actual level bytes.
+    int expected_level_bytes = CalculateLevelBytesLength(level_bytes, encoding_attributes);
+    if (static_cast<size_t>(expected_level_bytes) != level_bytes.size()) {
+        throw InvalidInputException("Level bytes size does not match encoding attributes");
+    }
+
+    if (page_type == "DATA_PAGE_V1") {
+        auto joined = Join(level_bytes, value_bytes);
+        return Compress(joined, compression);
+    }
+
+    if (page_type == "DATA_PAGE_V2") {
+        bool page_v2_is_compressed =
+            std::get<bool>(encoding_attributes.at("page_v2_is_compressed"));
+        if (page_v2_is_compressed) {
+            auto compressed_values = Compress(value_bytes, compression);
+            return Join(level_bytes, compressed_values);
+        } else {
+            return Join(level_bytes, value_bytes);
+        }
+    }
+
+    // DICTIONARY_PAGE has no level bytes.
+    if (page_type == "DICTIONARY_PAGE") {
+        return Compress(value_bytes, compression);
+    }
+
+    throw InvalidInputException("Unexpected page type: " + page_type);
 }
 
 TypedListValues ParseValueBytesIntoTypedList(

src/server/parquet_utils.h

@@ -77,13 +77,14 @@ LevelAndValueBytes DecompressAndSplit(
     const AttributesMap& encoding_attributes);
 
 /**
- * Merges level and value bytes and compresses them into plaintext.
- * Handles different page types (DATA_PAGE_V1, DATA_PAGE_V2, DICTIONARY_PAGE) appropriately.
- * Returns the joined and compressed plaintext.
+ * Reverse of DecompressAndSplit: joins level/value bytes and applies compression
+ * based on page type and encoding attributes.
  */
 std::vector<uint8_t> CompressAndJoin(
     const std::vector<uint8_t>& level_bytes,
-    const std::vector<uint8_t>& value_bytes);
+    const std::vector<uint8_t>& value_bytes,
+    CompressionCodec::type compression,
+    const AttributesMap& encoding_attributes);
 
 /**
  * Parse the value bytes into a typed list based on the data type and format.

src/server/parquet_utils_test.cpp

@@ -274,6 +274,13 @@ TEST(ParquetUtils, SliceValueBytesIntoRawBytes_FixedSizeMisaligned) {
         InvalidInputException);
 }
 
+TEST(ParquetUtils, SliceValueBytesIntoRawBytes_UnsupportedFormat) {
+    std::vector<uint8_t> bytes = {0x01,0x00,0x00,0x00};
+    EXPECT_THROW(
+        SliceValueBytesIntoRawBytes(bytes, Type::INT32, std::nullopt, Format::RLE),
+        DBPSUnsupportedException);
+}
+
 TEST(ParquetUtils, CombineRawBytesIntoValueBytes_INT32) {
     std::vector<RawValueBytes> elems = {
         {0x04,0x03,0x02,0x01},
@@ -309,6 +316,15 @@ TEST(ParquetUtils, CombineRawBytesIntoValueBytes_FIXED_LEN_BYTE_ARRAY_SizeMismat
         InvalidInputException);
 }
 
+TEST(ParquetUtils, CombineRawBytesIntoValueBytes_UnsupportedFormat) {
+    std::vector<RawValueBytes> elems = {
+        {0x04,0x03,0x02,0x01}
+    };
+    EXPECT_THROW(
+        CombineRawBytesIntoValueBytes(elems, Type::INT32, std::nullopt, Format::RLE),
+        DBPSUnsupportedException);
+}
+
 TEST(ParquetUtils, SliceAndCombine_RoundTrip_INT64) {
     // Two int64 values: little-endian bytes
     std::vector<uint8_t> bytes = {
@@ -402,3 +418,149 @@ TEST(ParquetUtils, DecompressAndSplit_DataPageV2_UnsupportedCompression) {
         DecompressAndSplit(plaintext, CompressionCodec::LZO, attribs),
         DBPSUnsupportedException);
 }
+
+TEST(ParquetUtils, CompressAndJoin_DataPageV1_Compressed) {
+    AttributesMap attribs = {
+        {"page_type", std::string("DATA_PAGE_V1")},
+        {"data_page_num_values", int32_t(2)},
+        {"data_page_max_repetition_level", int32_t(0)},
+        {"data_page_max_definition_level", int32_t(1)},
+        {"page_v1_repetition_level_encoding", std::string("RLE")},
+        {"page_v1_definition_level_encoding", std::string("RLE")}
+    };
+
+    std::vector<uint8_t> level_bytes;
+    append_u32_le(level_bytes, 2); // RLE block length
+    level_bytes.insert(level_bytes.end(), {0x0A, 0x0B});
+    std::vector<uint8_t> value_bytes = {0x21, 0x22, 0x23, 0x24};
+
+    auto joined = CompressAndJoin(level_bytes, value_bytes, CompressionCodec::SNAPPY, attribs);
+    auto expected = Compress(Join(level_bytes, value_bytes), CompressionCodec::SNAPPY);
+    EXPECT_EQ(joined, expected);
+
+    auto decomposed = DecompressAndSplit(joined, CompressionCodec::SNAPPY, attribs);
+    EXPECT_EQ(decomposed.level_bytes, level_bytes);
+    EXPECT_EQ(decomposed.value_bytes, value_bytes);
+}
+
+TEST(ParquetUtils, CompressAndJoin_DataPageV2_Uncompressed) {
+    AttributesMap attribs = {
+        {"page_type", std::string("DATA_PAGE_V2")},
+        {"data_page_num_values", int32_t(3)},
+        {"data_page_max_definition_level", int32_t(1)},
+        {"data_page_max_repetition_level", int32_t(0)},
+        {"page_v2_definition_levels_byte_length", int32_t(2)},
+        {"page_v2_repetition_levels_byte_length", int32_t(1)},
+        {"page_v2_num_nulls", int32_t(0)},
+        {"page_v2_is_compressed", false}
+    };
+
+    std::vector<uint8_t> level_bytes = {0x10, 0x11, 0x12}; // 2+1
+    std::vector<uint8_t> value_bytes = {0x21, 0x22, 0x23, 0x24};
+
+    auto joined = CompressAndJoin(level_bytes, value_bytes, CompressionCodec::UNCOMPRESSED, attribs);
+    std::vector<uint8_t> expected = Join(level_bytes, value_bytes);
+    EXPECT_EQ(joined, expected);
+
+    auto decomposed = DecompressAndSplit(joined, CompressionCodec::UNCOMPRESSED, attribs);
+    EXPECT_EQ(decomposed.level_bytes, level_bytes);
+    EXPECT_EQ(decomposed.value_bytes, value_bytes);
+}
+
+TEST(ParquetUtils, CompressAndJoin_DataPageV2_Compressed_RoundTrip) {
+    AttributesMap attribs = {
+        {"page_type", std::string("DATA_PAGE_V2")},
+        {"data_page_num_values", int32_t(4)},
+        {"data_page_max_definition_level", int32_t(1)},
+        {"data_page_max_repetition_level", int32_t(0)},
+        {"page_v2_definition_levels_byte_length", int32_t(2)},
+        {"page_v2_repetition_levels_byte_length", int32_t(1)},
+        {"page_v2_num_nulls", int32_t(0)},
+        {"page_v2_is_compressed", true}
+    };
+
+    std::vector<uint8_t> level_bytes = {0x10, 0x11, 0x12}; // len matches 2+1
+    std::vector<uint8_t> value_bytes = {0x21, 0x22, 0x23, 0x24, 0x25};
+
+    auto joined = CompressAndJoin(level_bytes, value_bytes, CompressionCodec::SNAPPY, attribs);
+    auto expected_compressed = Compress(value_bytes, CompressionCodec::SNAPPY);
+    auto expected_joined = Join(level_bytes, expected_compressed);
+    EXPECT_EQ(joined, expected_joined);
+
+    auto decomposed = DecompressAndSplit(joined, CompressionCodec::SNAPPY, attribs);
+
+    EXPECT_EQ(decomposed.level_bytes, level_bytes);
+    EXPECT_EQ(decomposed.value_bytes, value_bytes);
+}
+
+TEST(ParquetUtils, CompressAndJoin_DataPageV2_LevelLengthMismatch) {
+    AttributesMap attribs = {
+        {"page_type", std::string("DATA_PAGE_V2")},
+        {"data_page_num_values", int32_t(2)},
+        {"data_page_max_definition_level", int32_t(1)},
+        {"data_page_max_repetition_level", int32_t(0)},
+        {"page_v2_definition_levels_byte_length", int32_t(2)},
+        {"page_v2_repetition_levels_byte_length", int32_t(1)},
+        {"page_v2_num_nulls", int32_t(0)},
+        {"page_v2_is_compressed", true}
+    };
+
+    std::vector<uint8_t> level_bytes = {0x10, 0x11}; // expected 3 bytes -> mismatch
+    std::vector<uint8_t> value_bytes = {0x21, 0x22, 0x23};
+
+    EXPECT_THROW(
+        CompressAndJoin(level_bytes, value_bytes, CompressionCodec::SNAPPY, attribs),
+        InvalidInputException);
+}
+
+TEST(ParquetUtils, CompressAndJoin_DictionaryPage) {
+    AttributesMap attribs = {
+        {"page_type", std::string("DICTIONARY_PAGE")}
+    };
+
+    std::vector<uint8_t> level_bytes; // must be empty
+    std::vector<uint8_t> value_bytes = {0x31, 0x32, 0x33};
+
+    auto joined = CompressAndJoin(level_bytes, value_bytes, CompressionCodec::UNCOMPRESSED, attribs);
+    EXPECT_EQ(joined, value_bytes);
+}
+
+TEST(ParquetUtils, CompressAndJoin_UnsupportedCompression) {
+    AttributesMap attribs = {
+        {"page_type", std::string("DATA_PAGE_V1")},
+        {"data_page_num_values", int32_t(1)},
+        {"data_page_max_repetition_level", int32_t(0)},
+        {"data_page_max_definition_level", int32_t(0)},
+        {"page_v1_repetition_level_encoding", std::string("RLE")},
+        {"page_v1_definition_level_encoding", std::string("RLE")}
+    };
+
+    std::vector<uint8_t> level_bytes;
+    std::vector<uint8_t> value_bytes = {0x40};
+
+    EXPECT_THROW(
+        CompressAndJoin(level_bytes, value_bytes, CompressionCodec::LZO, attribs),
+        DBPSUnsupportedException);
+}
+
+TEST(ParquetUtils, CompressAndJoin_UnsupportedEncoding) {
+    AttributesMap attribs = {
+        {"page_type", std::string("DATA_PAGE_V1")},
+        {"data_page_num_values", int32_t(1)},
+        {"data_page_max_repetition_level", int32_t(1)}, // triggers repetition level parsing
+        {"data_page_max_definition_level", int32_t(0)},
+        {"page_v1_repetition_level_encoding", std::string("BIT_PACKED")}, // unsupported
+        {"page_v1_definition_level_encoding", std::string("RLE")}
+    };
+
+    // Build minimal level bytes that would be valid if encoding were RLE: len + payload
+    std::vector<uint8_t> level_bytes;
+    append_u32_le(level_bytes, 1);
+    level_bytes.push_back(0x00);
+
+    std::vector<uint8_t> value_bytes = {0x01, 0x02};
+
+    EXPECT_THROW(
+        CompressAndJoin(level_bytes, value_bytes, CompressionCodec::UNCOMPRESSED, attribs),
+        InvalidInputException);
+}