Add compression cache (#64)

* Implement caching for compression

* Handle cache type and some improvements

* Update more docstrings

* Add tests and some improvements

* Adapt stream_file_serializer and test

* Use local cache and add note

* Wrap compression cache and use unordered_map

* Use tuple as key to cache and add TODO note

Author: Hind-M

include/sparrow_ipc/chunk_memory_serializer.hpp

@@ -165,7 +165,8 @@ namespace sparrow_ipc
             std::vector<uint8_t> buffer;
             memory_output_stream stream(buffer);
             any_output_stream astream(stream);
-            serialize_record_batch(rb, astream, m_compression);
+            CompressionCache compressed_buffers_cache;
+            serialize_record_batch(rb, astream, m_compression, compressed_buffers_cache);
             m_pstream->write(std::move(buffer));
         }
     }

include/sparrow_ipc/compression.hpp

@@ -1,6 +1,7 @@
 #pragma once
 
 #include <cstdint>
+#include <optional>
 #include <span>
 #include <variant>
 #include <vector>
@@ -15,9 +16,41 @@ namespace sparrow_ipc
         ZSTD
     };
 
-    [[nodiscard]] SPARROW_IPC_API std::vector<std::uint8_t> compress(
+    class CompressionCacheImpl;
+
+    class SPARROW_IPC_API CompressionCache
+    {
+        public:
+            CompressionCache();
+            ~CompressionCache();
+
+            CompressionCache(CompressionCache&&) noexcept;
+            CompressionCache& operator=(CompressionCache&&) noexcept;
+
+            CompressionCache(const CompressionCache&) = delete;
+            CompressionCache& operator=(const CompressionCache&) = delete;
+
+            std::optional<std::span<const std::uint8_t>> find(const void* data_ptr, const size_t data_size);
+            std::span<const std::uint8_t> store(const void* data_ptr, const size_t data_size, std::vector<std::uint8_t>&& data);
+
+            size_t size() const;
+            size_t count(const void* data_ptr, const size_t data_size) const;
+            bool empty() const;
+            void clear();
+
+        private:
+            std::unique_ptr<CompressionCacheImpl> m_pimpl;
+    };
+
+    [[nodiscard]] SPARROW_IPC_API std::span<const std::uint8_t> compress(
         const CompressionType compression_type,
-        std::span<const std::uint8_t> data);
+        const std::span<const std::uint8_t>& data,
+        CompressionCache& cache);
+
+    [[nodiscard]] SPARROW_IPC_API size_t get_compressed_size(
+        const CompressionType compression_type,
+        const std::span<const std::uint8_t>& data,
+        CompressionCache& cache);
 
     [[nodiscard]] SPARROW_IPC_API std::variant<std::vector<std::uint8_t>, std::span<const std::uint8_t>> decompress(
         const CompressionType compression_type,

include/sparrow_ipc/deserialize.hpp

@@ -1,7 +1,7 @@
 #pragma once
 
 #include <optional>
-#include <string>
+#include <span>
 #include <vector>
 
 #include <sparrow/record_batch.hpp>
@@ -63,4 +63,4 @@ namespace sparrow_ipc
      */
     [[nodiscard]] SPARROW_IPC_API std::vector<sparrow::record_batch>
     deserialize_stream(std::span<const uint8_t> data);
-}
+}

include/sparrow_ipc/flatbuffer_utils.hpp

@@ -255,12 +255,14 @@ namespace sparrow_ipc
      * @param flatbuf_compressed_buffers A vector to store the resulting compressed buffer metadata.
      * @param offset The current offset in the buffer layout, which will be updated by the function.
      * @param compression_type The compression algorithm to use.
+     * @param cache A cache to store compressed buffers and avoid recompression.
      */
     void fill_compressed_buffers(
         const sparrow::arrow_proxy& arrow_proxy,
         std::vector<org::apache::arrow::flatbuf::Buffer>& flatbuf_compressed_buffers,
         int64_t& offset,
-        const CompressionType compression_type
+        const CompressionType compression_type,
+        CompressionCache& cache
     );
 
     /**
@@ -273,11 +275,14 @@ namespace sparrow_ipc
      *
      * @param record_batch The record batch whose buffers' compressed metadata is to be retrieved.
      * @param compression_type The compression algorithm that would be applied (e.g., LZ4_FRAME, ZSTD).
+     * @param cache A cache to store compressed buffers and avoid recompression.
      * @return A vector of FlatBuffer Buffer objects, each describing the offset and
      *         size of a corresponding compressed buffer within a larger message body.
      */
     [[nodiscard]] std::vector<org::apache::arrow::flatbuf::Buffer>
-    get_compressed_buffers(const sparrow::record_batch& record_batch, const CompressionType compression_type);
+    get_compressed_buffers(const sparrow::record_batch& record_batch,
+                           const CompressionType compression_type,
+                           CompressionCache& cache);
 
     /**
      * @brief Calculates the total aligned size in bytes of all buffers in an Arrow array structure.
@@ -286,11 +291,16 @@ namespace sparrow_ipc
      * in an Arrow array structure, including buffers from child arrays. Each
      * buffer size is aligned to 8-byte boundaries as required by the Arrow format.
      *
-     * @param arrow_proxy The Arrow array proxy containing buffers and child arrays
-     * @param compression The compression type to use when serializing
-     * @return int64_t The total aligned size in bytes of all buffers in the array hierarchy
+     * @param arrow_proxy The Arrow array proxy containing buffers and child arrays.
+     * @param compression Optional: The compression type to use when serializing.
+     * @param cache Optional: A cache to store and retrieve compressed buffer sizes, avoiding recompression.
+     * If compression is given, cache should be set as well.
+     * @return int64_t The total aligned size in bytes of all buffers in the array hierarchy.
+     * @throws std::invalid_argument if compression is given but not cache.
      */
-    [[nodiscard]] int64_t calculate_body_size(const sparrow::arrow_proxy& arrow_proxy, std::optional<CompressionType> compression = std::nullopt);
+    [[nodiscard]] int64_t calculate_body_size(const sparrow::arrow_proxy& arrow_proxy,
+                                            std::optional<CompressionType> compression = std::nullopt,
+                                            std::optional<std::reference_wrapper<CompressionCache>> cache = std::nullopt);
 
     /**
      * @brief Calculates the total body size of a record batch by summing the body sizes of all its columns.
@@ -299,11 +309,15 @@ namespace sparrow_ipc
      * the body size of each column's underlying Arrow array proxy. The body size represents
      * the total memory required for the serialized data content of the record batch.
      *
-     * @param record_batch The sparrow record batch containing columns to calculate size for
-     * @param compression The compression type to use when serializing
-     * @return int64_t The total body size in bytes of all columns in the record batch
+     * @param record_batch The sparrow record batch containing columns to calculate size for.
+     * @param compression Optional: The compression type to use when serializing. If not provided, sizes are for uncompressed buffers.
+     * @param cache Optional: A cache to store and retrieve compressed buffer sizes, avoiding recompression.
+     * If compression is given, cache should be set as well.
+     * @return int64_t The total body size in bytes of all columns in the record batch.
      */
-    [[nodiscard]] int64_t calculate_body_size(const sparrow::record_batch& record_batch, std::optional<CompressionType> compression = std::nullopt);
+    [[nodiscard]] int64_t calculate_body_size(const sparrow::record_batch& record_batch,
+                                            std::optional<CompressionType> compression = std::nullopt,
+                                            std::optional<std::reference_wrapper<CompressionCache>> cache = std::nullopt);
 
     /**
      * @brief Creates a FlatBuffer message containing a serialized Apache Arrow RecordBatch.
@@ -312,15 +326,20 @@ namespace sparrow_ipc
      * along with its metadata (field nodes and buffer information) into a FlatBuffer
      * format that conforms to the Arrow IPC specification.
      *
-     * @param record_batch The source record batch containing the data to be serialized
-     * @param compression Optional: The compression algorithm to be used for the message body
+     * @param record_batch The source record batch containing the data to be serialized.
+     * @param compression Optional: The compression algorithm to be used for the message body.
+     * @param cache Optional: A cache for compressed buffers to avoid recompression if compression is enabled.
+     * If compression is given, cache should be set as well.
      * @return A FlatBufferBuilder containing the complete serialized message ready for
      *         transmission or storage. The builder is finished and ready to be accessed
      *         via GetBufferPointer() and GetSize().
+     * @throws std::invalid_argument if compression is given but not cache.
      *
-     * @note The returned message uses Arrow IPC format version V5
-     * @note Variadic buffer counts is not currently implemented (set to 0)
+     * @note The returned message uses Arrow IPC format version V5.
+     * @note Variadic buffer counts is not currently implemented (set to 0).
      */
     [[nodiscard]] flatbuffers::FlatBufferBuilder
-    get_record_batch_message_builder(const sparrow::record_batch& record_batch, std::optional<CompressionType> compression = std::nullopt);
+    get_record_batch_message_builder(const sparrow::record_batch& record_batch,
+                                     std::optional<CompressionType> compression = std::nullopt,
+                                     std::optional<std::reference_wrapper<CompressionCache>> cache = std::nullopt);
 }

include/sparrow_ipc/serialize.hpp

@@ -27,8 +27,9 @@ namespace sparrow_ipc
      * @param record_batches Collection of record batches to serialize. All batches must have identical
      * schemas.
      * @param stream The output stream where the serialized data will be written.
-     * @param compression The compression type to use when serializing.
-     *
+     * @param compression Optional: The compression type to use when serializing.
+     * @param cache Optional: A cache to store and retrieve compressed buffers, avoiding recompression.
+     * If compression is given, cache should be set as well.
      * @throws std::invalid_argument If record batches have inconsistent schemas or if the collection
      *                               contains batches that cannot be serialized together.
      *
@@ -37,7 +38,9 @@ namespace sparrow_ipc
      */
     template <std::ranges::input_range R>
         requires std::same_as<std::ranges::range_value_t<R>, sparrow::record_batch>
-    void serialize_record_batches_to_ipc_stream(const R& record_batches, any_output_stream& stream, std::optional<CompressionType> compression)
+    void serialize_record_batches_to_ipc_stream(const R& record_batches, any_output_stream& stream,
+                                                std::optional<CompressionType> compression,
+                                                std::optional<std::reference_wrapper<CompressionCache>> cache)
     {
         if (record_batches.empty())
         {
@@ -53,7 +56,7 @@ namespace sparrow_ipc
         serialize_schema_message(record_batches[0], stream);
         for (const auto& rb : record_batches)
         {
-            serialize_record_batch(rb, stream, compression);
+            serialize_record_batch(rb, stream, compression, cache);
         }
         stream.write(end_of_stream);
     }
@@ -70,14 +73,18 @@ namespace sparrow_ipc
      *
      * @param record_batch The sparrow record batch to serialize
      * @param stream The output stream where the serialized record batch will be written
-     * @param compression The compression type to use when serializing.
-     *
+     * @param compression Optional: The compression type to use when serializing.
+     * @param cache Optional: A cache to store and retrieve compressed buffers, avoiding recompression.
+     * @note If compression is given, cache should be set as well.
      * @note The output follows Arrow IPC message format with proper alignment and
      *       includes both metadata and data portions of the record batch
      */
 
     SPARROW_IPC_API void
-    serialize_record_batch(const sparrow::record_batch& record_batch, any_output_stream& stream, std::optional<CompressionType> compression);
+    serialize_record_batch(const sparrow::record_batch& record_batch,
+                           any_output_stream& stream,
+                           std::optional<CompressionType> compression,
+                           std::optional<std::reference_wrapper<CompressionCache>> cache);
 
     /**
      * @brief Serializes a schema message for a record batch into a byte buffer.

include/sparrow_ipc/serialize_utils.hpp

@@ -39,12 +39,16 @@ namespace sparrow_ipc
      * The serialization follows the Arrow IPC stream format where each record batch message
      * consists of a metadata section followed by a body section containing the actual data.
      *
-     * @param record_batch The sparrow record batch to be serialized
-     * @param stream The output stream where the serialized record batch will be written
-     * @param compression The compression type to use when serializing
+     * @param record_batch The sparrow record batch to be serialized.
+     * @param stream The output stream where the serialized record batch will be written.
+     * @param compression Optional: The compression type to use when serializing.
+     * @param cache Optional: A cache for compressed buffers to avoid recompression if compression is enabled.
+     * If compression is given, cache should be set as well.
      */
     SPARROW_IPC_API void
-    serialize_record_batch(const sparrow::record_batch& record_batch, any_output_stream& stream, std::optional<CompressionType> compression);
+    serialize_record_batch(const sparrow::record_batch& record_batch, any_output_stream& stream,
+                           std::optional<CompressionType> compression,
+                           std::optional<std::reference_wrapper<CompressionCache>> cache);
 
     /**
      * @brief Calculates the total serialized size of a schema message.
@@ -73,28 +77,36 @@ namespace sparrow_ipc
      * - Padding to 8-byte alignment after metadata
      * - Body data with 8-byte alignment between buffers
      *
-     * @param record_batch The record batch to be measured
-     * @param compression The compression type to use when serializing
-     * @return The total size in bytes that the serialized record batch would occupy
+     * @param record_batch The record batch to be measured.
+     * @param compression Optional: The compression type to use when serializing.
+     * @param cache Optional: A cache to store and retrieve compressed buffer sizes, avoiding recompression.
+     * If compression is given, cache should be set as well.
+     * @return The total size in bytes that the serialized record batch would occupy.
      */
     [[nodiscard]] SPARROW_IPC_API std::size_t
-    calculate_record_batch_message_size(const sparrow::record_batch& record_batch, std::optional<CompressionType> compression = std::nullopt);
+    calculate_record_batch_message_size(const sparrow::record_batch& record_batch,
+                                        std::optional<CompressionType> compression = std::nullopt,
+                                        std::optional<std::reference_wrapper<CompressionCache>> cache = std::nullopt);
 
     /**
      * @brief Calculates the total serialized size for a collection of record batches.
      *
      * This function computes the complete size that would be produced by serializing
      * a schema message followed by all record batch messages in the collection.
      *
-     * @tparam R Range type containing sparrow::record_batch objects
-     * @param record_batches Collection of record batches to be measured
-     * @param compression The compression type to use when serializing
-     * @return The total size in bytes for the complete serialized output
-     * @throws std::invalid_argument if record batches have inconsistent schemas
+     * @tparam R Range type containing sparrow::record_batch objects.
+     * @param record_batches Collection of record batches to be measured.
+     * @param compression Optional: The compression type to use when serializing.
+     * @param cache Optional: A cache to store and retrieve compressed buffer sizes, avoiding recompression.
+     * If compression is given, cache should be set as well.
+     * @return The total size in bytes for the complete serialized output.
+     * @throws std::invalid_argument if record batches have inconsistent schemas.
      */
     template <std::ranges::input_range R>
         requires std::same_as<std::ranges::range_value_t<R>, sparrow::record_batch>
-    [[nodiscard]] std::size_t calculate_total_serialized_size(const R& record_batches, std::optional<CompressionType> compression = std::nullopt)
+    [[nodiscard]] std::size_t calculate_total_serialized_size(const R& record_batches,
+                                                              std::optional<CompressionType> compression = std::nullopt,
+                                                              std::optional<std::reference_wrapper<CompressionCache>> cache = std::nullopt)
     {
         if (record_batches.empty())
         {
@@ -113,7 +125,7 @@ namespace sparrow_ipc
         // Calculate record batch message sizes
         for (const auto& record_batch : record_batches)
         {
-            total_size += calculate_record_batch_message_size(record_batch, compression);
+            total_size += calculate_record_batch_message_size(record_batch, compression, cache);
         }
 
         return total_size;
@@ -131,11 +143,16 @@ namespace sparrow_ipc
      * 8-byte boundary, which is typically required for efficient memory access and Arrow
      * format compliance.
      *
-     * @param arrow_proxy The arrow proxy containing buffers and potential child proxies to serialize
-     * @param stream The output stream where the serialized body data will be written
-     * @param compression The compression type to use when serializing
+     * @param arrow_proxy The arrow proxy containing buffers and potential child proxies to serialize.
+     * @param stream The output stream where the serialized body data will be written.
+     * @param compression Optional: The compression type to use when serializing.
+     * @param cache Optional: A cache for compressed buffers to avoid recompression if compression is enabled.
+     * If compression is given, cache should be set as well.
+     * @throws std::invalid_argument if compression is given but not cache.
      */
-    SPARROW_IPC_API void fill_body(const sparrow::arrow_proxy& arrow_proxy, any_output_stream& stream, std::optional<CompressionType> compression = std::nullopt);
+    SPARROW_IPC_API void fill_body(const sparrow::arrow_proxy& arrow_proxy, any_output_stream& stream,
+                                   std::optional<CompressionType> compression = std::nullopt,
+                                   std::optional<std::reference_wrapper<CompressionCache>> cache = std::nullopt);
 
     /**
      * @brief Generates a serialized body from a record batch.
@@ -144,11 +161,15 @@ namespace sparrow_ipc
      * extracts their Arrow proxy representations, and serializes them into a
      * single byte vector that forms the body of the serialized data.
      *
-     * @param record_batch The record batch containing columns to be serialized
-     * @param stream The output stream where the serialized body will be written
-     * @param compression The compression type to use when serializing
+     * @param record_batch The record batch containing columns to be serialized.
+     * @param stream The output stream where the serialized body will be written.
+     * @param compression Optional: The compression type to use when serializing.
+     * @param cache Optional: A cache for compressed buffers to avoid recompression if compression is enabled.
+     * If compression is given, cache should be set as well.
      */
-    SPARROW_IPC_API void generate_body(const sparrow::record_batch& record_batch, any_output_stream& stream, std::optional<CompressionType> compression = std::nullopt);
+    SPARROW_IPC_API void generate_body(const sparrow::record_batch& record_batch, any_output_stream& stream,
+                                       std::optional<CompressionType> compression = std::nullopt,
+                                       std::optional<std::reference_wrapper<CompressionCache>> cache = std::nullopt);
 
     SPARROW_IPC_API std::vector<sparrow::data_type> get_column_dtypes(const sparrow::record_batch& rb);
 }