Obtain compressed size and number of elements from page header

Author: sfc-gh-pgaur

cpp/src/arrow/util/alp/alp_wrapper.cc

@@ -37,32 +37,28 @@ namespace {
 
 /// \brief Header structure for ALP compression blocks
 ///
-/// Contains metadata required to decompress the data.
+/// Contains metadata required to decompress the data. Note that compressed_size
+/// and num_elements are NOT stored in the header - they are available from the
+/// page header and passed to the Decode() function.
 ///
 /// Serialization format (version 1):
 ///
 ///   +---------------------------------------------------+
-///   |  CompressionBlockHeader (40 bytes)                |
+///   |  CompressionBlockHeader (24 bytes)                |
 ///   +---------------------------------------------------+
 ///   |  Offset |  Field              |  Size             |
 ///   +---------+---------------------+-------------------+
 ///   |    0    |  version            |  8 bytes (uint64) |
-///   |    8    |  compressed_size    |  8 bytes (uint64) |
-///   |   16    |  num_elements       |  8 bytes (uint64) |
-///   |   24    |  vector_size        |  8 bytes (uint64) |
-///   |   32    |  compression_mode   |  4 bytes (enum)   |
-///   |   36    |  bit_pack_layout    |  4 bytes (enum)   |
+///   |    8    |  vector_size        |  8 bytes (uint64) |
+///   |   16    |  compression_mode   |  4 bytes (enum)   |
+///   |   20    |  bit_pack_layout    |  4 bytes (enum)   |
 ///   +---------------------------------------------------+
 ///
 /// \note version must remain the first field to allow reading the rest
 ///       of the header based on version number.
 struct CompressionBlockHeader {
   /// Version number. Must remain the first field for version-based parsing.
   uint64_t version = 0;
-  /// Size of the compressed data in bytes (includes header).
-  uint64_t compressed_size = 0;
-  /// Number of elements in the compressed data.
-  uint64_t num_elements = 0;
   /// Vector size used for compression.
   /// Must be AlpConstants::kAlpVectorSize for decompression.
   uint64_t vector_size = 0;
@@ -78,8 +74,8 @@ struct CompressionBlockHeader {
   static size_t GetSizeForVersion(uint64_t v) {
     size_t size;
     if (v == 1) {
-      size = sizeof(version) + sizeof(compressed_size) + sizeof(num_elements) +
-             sizeof(vector_size) + sizeof(compression_mode) + sizeof(bit_pack_layout);
+      size = sizeof(version) + sizeof(vector_size) + sizeof(compression_mode) +
+             sizeof(bit_pack_layout);
     } else {
       ARROW_CHECK(false) << "unknown_version: " << v;
     }
@@ -151,33 +147,24 @@ void AlpWrapper<T>::Encode(const T* decomp, size_t decomp_size, char* comp,
 
   CompressionBlockHeader header{};
   header.version = version;
-  header.compressed_size =
-      ::arrow::util::alp::CompressionBlockHeader::GetSizeForVersion(version) +
-      compression_progress.num_compressed_bytes_produced;
-  header.num_elements = decomp_size / sizeof(T);
   header.vector_size = AlpConstants::kAlpVectorSize;
   header.compression_mode = AlpMode::kAlp;
   header.bit_pack_layout = AlpBitPackLayout::kNormal;
 
   std::memcpy(encoded_header, &header,
               ::arrow::util::alp::CompressionBlockHeader::GetSizeForVersion(version));
-  *comp_size = header.compressed_size;
+  *comp_size = ::arrow::util::alp::CompressionBlockHeader::GetSizeForVersion(version) +
+               compression_progress.num_compressed_bytes_produced;
 }
 
 template <typename T>
 template <typename TargetType>
-void AlpWrapper<T>::Decode(TargetType* decomp, size_t* decomp_size, const char* comp,
+void AlpWrapper<T>::Decode(TargetType* decomp, uint64_t num_elements, const char* comp,
                            size_t comp_size) {
   const CompressionBlockHeader header = LoadHeader(comp, comp_size);
   ARROW_CHECK(header.vector_size == AlpConstants::kAlpVectorSize)
       << "unsupported_vector_size: " << header.vector_size;
 
-  if (header.num_elements * sizeof(TargetType) > *decomp_size) {
-    *decomp_size = 0;
-    return;
-  }
-
-  const uint64_t elements_to_decode = header.num_elements;
   const char* compression_body =
       comp + ::arrow::util::alp::CompressionBlockHeader::GetSizeForVersion(header.version);
   const uint64_t compression_body_size =
@@ -186,34 +173,17 @@ void AlpWrapper<T>::Decode(TargetType* decomp, size_t* decomp_size, const char*
 
   ARROW_CHECK(header.compression_mode == AlpMode::kAlp) << "alp_decode_unsupported_mode";
 
-  uint64_t elements_decoded =
-      DecodeAlp(decomp, elements_to_decode, compression_body, compression_body_size,
-                header.bit_pack_layout)
-          .num_decompressed_elements_produced;
-  *decomp_size = elements_decoded * sizeof(TargetType);
+  DecodeAlp<TargetType>(decomp, num_elements, compression_body, compression_body_size,
+                        header.bit_pack_layout);
 }
 
-template void AlpWrapper<float>::Decode(float* decomp, size_t* decomp_size,
+template void AlpWrapper<float>::Decode(float* decomp, uint64_t num_elements,
                                         const char* comp, size_t comp_size);
-template void AlpWrapper<float>::Decode(double* decomp, size_t* decomp_size,
+template void AlpWrapper<float>::Decode(double* decomp, uint64_t num_elements,
                                         const char* comp, size_t comp_size);
-template void AlpWrapper<double>::Decode(double* decomp, size_t* decomp_size,
+template void AlpWrapper<double>::Decode(double* decomp, uint64_t num_elements,
                                          const char* comp, size_t comp_size);
 
-template <typename T>
-template <typename TargetType>
-uint64_t AlpWrapper<T>::GetDecompressedSize(const char* comp, uint64_t comp_size) {
-  const CompressionBlockHeader header = LoadHeader(comp, comp_size);
-  return header.num_elements * sizeof(TargetType);
-}
-
-template uint64_t AlpWrapper<float>::GetDecompressedSize<float>(const char* comp,
-                                                                uint64_t comp_size);
-template uint64_t AlpWrapper<float>::GetDecompressedSize<double>(const char* comp,
-                                                                 uint64_t comp_size);
-template uint64_t AlpWrapper<double>::GetDecompressedSize<double>(const char* comp,
-                                                                  uint64_t comp_size);
-
 template <typename T>
 uint64_t AlpWrapper<T>::GetMaxCompressedSize(uint64_t decomp_size) {
   ARROW_CHECK(decomp_size % sizeof(T) == 0)

cpp/src/arrow/util/alp/alp_wrapper.h

@@ -63,30 +63,17 @@ class AlpWrapper {
   /// \brief Decode floating point values
   ///
   /// \param[out] decomp pointer to the memory region we will decode into.
-  ///             The caller is responsible for ensuring this is big enough.
-  /// \param[in,out] decomp_size the actual size of decoded data in bytes,
-  ///                expects the decomp size as input.
+  ///             The caller is responsible for ensuring this is big enough
+  ///             to hold num_elements values.
+  /// \param[in] num_elements number of elements to decode (from page header)
   /// \param[in] comp pointer to the input that is to be decoded
-  /// \param[in] comp_size size of the input in bytes.
+  /// \param[in] comp_size size of the input in bytes (from page header)
   /// \tparam TargetType the type that is used to store the output.
   ///         May not be a narrowing conversion from T.
   template <typename TargetType>
-  static void Decode(TargetType* decomp, size_t* decomp_size, const char* comp,
+  static void Decode(TargetType* decomp, uint64_t num_elements, const char* comp,
                      size_t comp_size);
 
-  /// \brief Get the decompressed size of a compression block
-  ///
-  /// Get the size of a compression block encoded previously with
-  /// AlpWrapper::Encode().
-  ///
-  /// \param[in] comp start of the memory region containing the compression block
-  /// \param[in] comp_size size of the compression block
-  /// \return the decompressed size of the block, in bytes
-  /// \tparam TargetType the type that is used to store the output.
-  ///         May not be a narrowing conversion from T.
-  template <typename TargetType>
-  static uint64_t GetDecompressedSize(const char* comp, uint64_t comp_size);
-
   /// \brief Get the maximum compressed size of an uncompressed buffer
   ///
   /// \param[in] decomp_size the size of the uncompressed buffer in bytes

cpp/src/parquet/encoding_alp_benchmark.cc

@@ -31,7 +31,7 @@
 #include <benchmark/benchmark.h>
 
 #include "arrow/buffer.h"
-#include "arrow/util/alp/AlpWrapper.h"
+#include "arrow/util/alp/alp_wrapper.h"
 #include "arrow/util/compression.h"
 #include "parquet/encoding.h"
 #include "parquet/schema.h"