blob: 971e6ccebc9d08a894c4eb74cbf0bc834c4e10f9 [file]
// Licensed to the Apache Software Foundation (ASF) under one
// or more contributor license agreements. See the NOTICE file
// distributed with this work for additional information
// regarding copyright ownership. The ASF licenses this file
// to you under the Apache License, Version 2.0 (the
// "License"); you may not use this file except in compliance
// with the License. You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing,
// software distributed under the License is distributed on an
// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
// KIND, either express or implied. See the License for the
// specific language governing permissions and limitations
// under the License.
#pragma once
#include "parquet/windows_compatibility.h"
#include <cstdint>
#include <limits>
#include <memory>
#include <span>
#include <sstream>
#include <string>
#include <type_traits>
#include <utility>
#include <vector>
// TCompactProtocol requires some #defines to work right.
#define SIGNED_RIGHT_SHIFT_IS 1
#define ARITHMETIC_RIGHT_SHIFT 1
#include <thrift/TApplicationException.h>
#include <thrift/protocol/TCompactProtocol.h>
#include <thrift/transport/TBufferTransports.h>
#include "arrow/util/logging.h"
#include "parquet/encryption/internal_file_decryptor.h"
#include "parquet/encryption/internal_file_encryptor.h"
#include "parquet/exception.h"
#include "parquet/geospatial/statistics.h"
#include "parquet/platform.h"
#include "parquet/properties.h"
#include "parquet/schema.h"
#include "parquet/size_statistics.h"
#include "parquet/statistics.h"
#include "parquet/types.h"
#include "generated/parquet_types.h" // IWYU pragma: export
namespace parquet {
// ----------------------------------------------------------------------
// Convert Thrift enums to Parquet enums
// Unsafe enum converters (input is not checked for validity)
constexpr Type::type FromThriftUnsafe(format::Type::type type) {
return static_cast<Type::type>(type);
}
constexpr ConvertedType::type FromThriftUnsafe(format::ConvertedType::type type) {
// item 0 is NONE
return static_cast<ConvertedType::type>(static_cast<int>(type) + 1);
}
constexpr Repetition::type FromThriftUnsafe(format::FieldRepetitionType::type type) {
return static_cast<Repetition::type>(type);
}
static inline Encoding::type FromThriftUnsafe(format::Encoding::type type) {
return static_cast<Encoding::type>(type);
}
constexpr PageType::type FromThriftUnsafe(format::PageType::type type) {
return static_cast<PageType::type>(type);
}
constexpr Compression::type FromThriftUnsafe(format::CompressionCodec::type type) {
switch (type) {
case format::CompressionCodec::UNCOMPRESSED:
return Compression::UNCOMPRESSED;
case format::CompressionCodec::SNAPPY:
return Compression::SNAPPY;
case format::CompressionCodec::GZIP:
return Compression::GZIP;
case format::CompressionCodec::LZO:
return Compression::LZO;
case format::CompressionCodec::BROTLI:
return Compression::BROTLI;
case format::CompressionCodec::LZ4:
return Compression::LZ4_HADOOP;
case format::CompressionCodec::LZ4_RAW:
return Compression::LZ4;
case format::CompressionCodec::ZSTD:
return Compression::ZSTD;
default:
ARROW_DCHECK(false) << "Cannot reach here";
return Compression::UNCOMPRESSED;
}
}
constexpr BoundaryOrder::type FromThriftUnsafe(format::BoundaryOrder::type type) {
return static_cast<BoundaryOrder::type>(type);
}
constexpr GeometryLogicalType::EdgeInterpolationAlgorithm FromThriftUnsafe(
format::EdgeInterpolationAlgorithm::type type) {
switch (type) {
case format::EdgeInterpolationAlgorithm::SPHERICAL:
return GeometryLogicalType::EdgeInterpolationAlgorithm::SPHERICAL;
case format::EdgeInterpolationAlgorithm::VINCENTY:
return GeometryLogicalType::EdgeInterpolationAlgorithm::VINCENTY;
case format::EdgeInterpolationAlgorithm::THOMAS:
return GeometryLogicalType::EdgeInterpolationAlgorithm::THOMAS;
case format::EdgeInterpolationAlgorithm::ANDOYER:
return GeometryLogicalType::EdgeInterpolationAlgorithm::ANDOYER;
case format::EdgeInterpolationAlgorithm::KARNEY:
return GeometryLogicalType::EdgeInterpolationAlgorithm::KARNEY;
default:
ARROW_DCHECK(false) << "Cannot reach here";
return GeometryLogicalType::EdgeInterpolationAlgorithm::UNKNOWN;
}
}
namespace internal {
template <typename T>
struct ThriftEnumTypeTraits {};
template <>
struct ThriftEnumTypeTraits<::parquet::format::Type::type> {
using ParquetEnum = Type;
};
template <>
struct ThriftEnumTypeTraits<::parquet::format::ConvertedType::type> {
using ParquetEnum = ConvertedType;
};
template <>
struct ThriftEnumTypeTraits<::parquet::format::FieldRepetitionType::type> {
using ParquetEnum = Repetition;
};
template <>
struct ThriftEnumTypeTraits<::parquet::format::Encoding::type> {
using ParquetEnum = Encoding;
};
template <>
struct ThriftEnumTypeTraits<::parquet::format::PageType::type> {
using ParquetEnum = PageType;
};
template <>
struct ThriftEnumTypeTraits<::parquet::format::BoundaryOrder::type> {
using ParquetEnum = BoundaryOrder;
};
// If the parquet file is corrupted it is possible the enum value decoded
// will not be in the range of defined values, which is undefined behaviour.
// This facility prevents this by loading the value as the underlying type
// and checking to make sure it is in range.
template <typename EnumType,
typename EnumTypeRaw = typename std::underlying_type<EnumType>::type>
inline static EnumTypeRaw LoadEnumRaw(const EnumType* in) {
EnumTypeRaw raw_value;
// Use memcpy(), as a regular cast would be undefined behaviour on invalid values
memcpy(&raw_value, in, sizeof(EnumType));
return raw_value;
}
template <typename ApiType>
struct SafeLoader {
using ApiTypeEnum = typename ApiType::type;
using ApiTypeRawEnum = typename std::underlying_type<ApiTypeEnum>::type;
template <typename ThriftType>
inline static ApiTypeRawEnum LoadRaw(const ThriftType* in) {
static_assert(sizeof(ApiTypeEnum) == sizeof(ThriftType),
"parquet type should always be the same size as thrift type");
return static_cast<ApiTypeRawEnum>(LoadEnumRaw(in));
}
template <typename ThriftType>
inline static ApiTypeEnum Load(const ThriftType* in) {
const auto raw_value = LoadRaw(in);
if constexpr (std::is_unsigned_v<ApiTypeRawEnum>) {
if (ARROW_PREDICT_FALSE(raw_value >=
static_cast<ApiTypeRawEnum>(ApiType::UNDEFINED))) {
return ApiType::UNDEFINED;
}
} else {
if (ARROW_PREDICT_FALSE(raw_value >=
static_cast<ApiTypeRawEnum>(ApiType::UNDEFINED) ||
raw_value < 0)) {
return ApiType::UNDEFINED;
}
}
return FromThriftUnsafe(static_cast<ThriftType>(raw_value));
}
};
} // namespace internal
// Safe enum loader: will check for invalid enum value before converting
template <typename ThriftType,
typename ParquetEnum =
typename internal::ThriftEnumTypeTraits<ThriftType>::ParquetEnum>
inline typename ParquetEnum::type LoadEnumSafe(const ThriftType* in) {
return internal::SafeLoader<ParquetEnum>::Load(in);
}
inline typename Compression::type LoadEnumSafe(const format::CompressionCodec::type* in) {
const auto raw_value = internal::LoadEnumRaw(in);
// Check bounds manually, as Compression::type doesn't have the same values
// as format::CompressionCodec.
const auto min_value =
static_cast<decltype(raw_value)>(format::CompressionCodec::UNCOMPRESSED);
const auto max_value =
static_cast<decltype(raw_value)>(format::CompressionCodec::LZ4_RAW);
if (raw_value < min_value || raw_value > max_value) {
return Compression::UNCOMPRESSED;
}
return FromThriftUnsafe(*in);
}
inline typename LogicalType::EdgeInterpolationAlgorithm LoadEnumSafe(
const format::EdgeInterpolationAlgorithm::type* in) {
const auto raw_value = internal::LoadEnumRaw(in);
if (ARROW_PREDICT_FALSE(raw_value < format::EdgeInterpolationAlgorithm::SPHERICAL ||
raw_value > format::EdgeInterpolationAlgorithm::KARNEY)) {
return LogicalType::EdgeInterpolationAlgorithm::UNKNOWN;
}
return FromThriftUnsafe(*in);
}
// Safe non-enum converters
static inline AadMetadata FromThrift(format::AesGcmV1 aesGcmV1) {
return AadMetadata{aesGcmV1.aad_prefix, aesGcmV1.aad_file_unique,
aesGcmV1.supply_aad_prefix};
}
static inline AadMetadata FromThrift(format::AesGcmCtrV1 aesGcmCtrV1) {
return AadMetadata{aesGcmCtrV1.aad_prefix, aesGcmCtrV1.aad_file_unique,
aesGcmCtrV1.supply_aad_prefix};
}
// Selects how thrift Statistics min/max fields should populate EncodedStatistics.
enum class StatisticsMinMaxField {
// Do not populate min/max, because the ordering is undefined or unsupported.
kInvalid,
// Populate min/max from the min_value/max_value fields.
kMinValueMaxValue,
// Populate min/max from the legacy min/max fields.
kLegacyMinMax,
};
// Keep this field-selection logic consistent with ColumnDescriptor::can_use_min_max().
static inline StatisticsMinMaxField GetStatisticsMinMaxField(
const ColumnDescriptor& descr) {
switch (descr.column_order().get_order()) {
case ColumnOrder::TYPE_DEFINED_ORDER:
return descr.sort_order() != SortOrder::UNKNOWN
? StatisticsMinMaxField::kMinValueMaxValue
: StatisticsMinMaxField::kInvalid;
case ColumnOrder::UNDEFINED:
return descr.sort_order() == SortOrder::SIGNED
? StatisticsMinMaxField::kLegacyMinMax
: StatisticsMinMaxField::kInvalid;
case ColumnOrder::UNKNOWN:
return StatisticsMinMaxField::kInvalid;
}
return StatisticsMinMaxField::kInvalid;
}
static inline EncodedStatistics FromThrift(const format::Statistics& stats,
StatisticsMinMaxField min_max) {
EncodedStatistics out;
if (min_max == StatisticsMinMaxField::kMinValueMaxValue) {
if (stats.__isset.max_value) {
out.set_max(stats.max_value);
if (stats.__isset.is_max_value_exact) {
out.is_max_value_exact = stats.is_max_value_exact;
}
}
if (stats.__isset.min_value) {
out.set_min(stats.min_value);
if (stats.__isset.is_min_value_exact) {
out.is_min_value_exact = stats.is_min_value_exact;
}
}
} else if (min_max == StatisticsMinMaxField::kLegacyMinMax) {
if (stats.__isset.max) {
out.set_max(stats.max);
}
if (stats.__isset.min) {
out.set_min(stats.min);
}
}
if (stats.__isset.null_count) {
out.set_null_count(stats.null_count);
}
if (stats.__isset.distinct_count) {
out.set_distinct_count(stats.distinct_count);
}
return out;
}
static inline geospatial::EncodedGeoStatistics FromThrift(
const format::GeospatialStatistics& geo_stats) {
geospatial::EncodedGeoStatistics out;
out.geospatial_types = geo_stats.geospatial_types;
if (geo_stats.__isset.bbox) {
out.xmin = geo_stats.bbox.xmin;
out.xmax = geo_stats.bbox.xmax;
out.ymin = geo_stats.bbox.ymin;
out.ymax = geo_stats.bbox.ymax;
out.xy_bounds_present = true;
if (geo_stats.bbox.__isset.zmin && geo_stats.bbox.__isset.zmax) {
out.zmin = geo_stats.bbox.zmin;
out.zmax = geo_stats.bbox.zmax;
out.z_bounds_present = true;
}
if (geo_stats.bbox.__isset.mmin && geo_stats.bbox.__isset.mmax) {
out.mmin = geo_stats.bbox.mmin;
out.mmax = geo_stats.bbox.mmax;
out.m_bounds_present = true;
}
}
return out;
}
static inline format::EdgeInterpolationAlgorithm::type ToThrift(
LogicalType::EdgeInterpolationAlgorithm algorithm) {
switch (algorithm) {
case LogicalType::EdgeInterpolationAlgorithm::SPHERICAL:
return format::EdgeInterpolationAlgorithm::SPHERICAL;
case LogicalType::EdgeInterpolationAlgorithm::VINCENTY:
return format::EdgeInterpolationAlgorithm::VINCENTY;
case LogicalType::EdgeInterpolationAlgorithm::THOMAS:
return format::EdgeInterpolationAlgorithm::THOMAS;
case LogicalType::EdgeInterpolationAlgorithm::ANDOYER:
return format::EdgeInterpolationAlgorithm::ANDOYER;
case LogicalType::EdgeInterpolationAlgorithm::KARNEY:
return format::EdgeInterpolationAlgorithm::KARNEY;
default:
throw ParquetException("Unknown value for geometry algorithm: ",
static_cast<int>(algorithm));
}
}
static inline EncryptionAlgorithm FromThrift(format::EncryptionAlgorithm encryption) {
EncryptionAlgorithm encryption_algorithm;
if (encryption.__isset.AES_GCM_V1) {
encryption_algorithm.algorithm = ParquetCipher::AES_GCM_V1;
encryption_algorithm.aad = FromThrift(encryption.AES_GCM_V1);
} else if (encryption.__isset.AES_GCM_CTR_V1) {
encryption_algorithm.algorithm = ParquetCipher::AES_GCM_CTR_V1;
encryption_algorithm.aad = FromThrift(encryption.AES_GCM_CTR_V1);
} else {
throw ParquetException("Unsupported algorithm");
}
return encryption_algorithm;
}
static inline SortingColumn FromThrift(format::SortingColumn thrift_sorting_column) {
SortingColumn sorting_column;
sorting_column.column_idx = thrift_sorting_column.column_idx;
sorting_column.nulls_first = thrift_sorting_column.nulls_first;
sorting_column.descending = thrift_sorting_column.descending;
return sorting_column;
}
static inline SizeStatistics FromThrift(const format::SizeStatistics& size_stats) {
return SizeStatistics{
size_stats.definition_level_histogram, size_stats.repetition_level_histogram,
size_stats.__isset.unencoded_byte_array_data_bytes
? std::make_optional(size_stats.unencoded_byte_array_data_bytes)
: std::nullopt};
}
// ----------------------------------------------------------------------
// Convert Thrift enums from Parquet enums
constexpr format::Type::type ToThrift(Type::type type) {
return static_cast<format::Type::type>(type);
}
constexpr format::ConvertedType::type ToThrift(ConvertedType::type type) {
// item 0 is NONE
ARROW_DCHECK_NE(type, ConvertedType::NONE);
// it is forbidden to emit "NA" (PARQUET-1990)
ARROW_DCHECK_NE(type, ConvertedType::NA);
ARROW_DCHECK_NE(type, ConvertedType::UNDEFINED);
return static_cast<format::ConvertedType::type>(static_cast<int>(type) - 1);
}
constexpr format::FieldRepetitionType::type ToThrift(Repetition::type type) {
return static_cast<format::FieldRepetitionType::type>(type);
}
constexpr format::Encoding::type ToThrift(Encoding::type type) {
return static_cast<format::Encoding::type>(type);
}
constexpr format::CompressionCodec::type ToThrift(Compression::type type) {
switch (type) {
case Compression::UNCOMPRESSED:
return format::CompressionCodec::UNCOMPRESSED;
case Compression::SNAPPY:
return format::CompressionCodec::SNAPPY;
case Compression::GZIP:
return format::CompressionCodec::GZIP;
case Compression::LZO:
return format::CompressionCodec::LZO;
case Compression::BROTLI:
return format::CompressionCodec::BROTLI;
case Compression::LZ4:
return format::CompressionCodec::LZ4_RAW;
case Compression::LZ4_HADOOP:
// Deprecated "LZ4" Parquet compression has Hadoop-specific framing
return format::CompressionCodec::LZ4;
case Compression::ZSTD:
return format::CompressionCodec::ZSTD;
default:
ARROW_DCHECK(false) << "Cannot reach here";
return format::CompressionCodec::UNCOMPRESSED;
}
}
constexpr format::BoundaryOrder::type ToThrift(BoundaryOrder::type type) {
switch (type) {
case BoundaryOrder::Unordered:
case BoundaryOrder::Ascending:
case BoundaryOrder::Descending:
return static_cast<format::BoundaryOrder::type>(type);
default:
ARROW_DCHECK(false) << "Cannot reach here";
return format::BoundaryOrder::UNORDERED;
}
}
static inline format::SortingColumn ToThrift(SortingColumn sorting_column) {
format::SortingColumn thrift_sorting_column;
thrift_sorting_column.column_idx = sorting_column.column_idx;
thrift_sorting_column.descending = sorting_column.descending;
thrift_sorting_column.nulls_first = sorting_column.nulls_first;
return thrift_sorting_column;
}
static inline format::GeospatialStatistics ToThrift(
const geospatial::EncodedGeoStatistics& encoded_geo_stats) {
format::GeospatialStatistics geospatial_statistics;
geospatial_statistics.__set_geospatial_types(encoded_geo_stats.geospatial_types);
if (encoded_geo_stats.xy_bounds_present) {
format::BoundingBox bbox;
bbox.__set_xmin(encoded_geo_stats.xmin);
bbox.__set_xmax(encoded_geo_stats.xmax);
bbox.__set_ymin(encoded_geo_stats.ymin);
bbox.__set_ymax(encoded_geo_stats.ymax);
if (encoded_geo_stats.z_bounds_present) {
bbox.__set_zmin(encoded_geo_stats.zmin);
bbox.__set_zmax(encoded_geo_stats.zmax);
}
if (encoded_geo_stats.m_bounds_present) {
bbox.__set_mmin(encoded_geo_stats.mmin);
bbox.__set_mmax(encoded_geo_stats.mmax);
}
geospatial_statistics.__set_bbox(std::move(bbox));
}
return geospatial_statistics;
}
static inline format::Statistics ToThrift(const EncodedStatistics& stats) {
format::Statistics statistics;
if (stats.has_min) {
statistics.__set_min_value(stats.min());
if (stats.is_min_value_exact.has_value()) {
statistics.__set_is_min_value_exact(stats.is_min_value_exact.value());
}
// If the order is SIGNED, then the old min value must be set too.
// This for backward compatibility
if (stats.is_signed()) {
statistics.__set_min(stats.min());
}
}
if (stats.has_max) {
statistics.__set_max_value(stats.max());
if (stats.is_max_value_exact.has_value()) {
statistics.__set_is_max_value_exact(stats.is_max_value_exact.value());
}
// If the order is SIGNED, then the old max value must be set too.
// This for backward compatibility
if (stats.is_signed()) {
statistics.__set_max(stats.max());
}
}
if (stats.has_null_count) {
statistics.__set_null_count(stats.null_count);
}
if (stats.has_distinct_count) {
statistics.__set_distinct_count(stats.distinct_count);
}
return statistics;
}
static inline format::AesGcmV1 ToAesGcmV1Thrift(AadMetadata aad) {
format::AesGcmV1 aesGcmV1;
// aad_file_unique is always set
aesGcmV1.__set_aad_file_unique(aad.aad_file_unique);
aesGcmV1.__set_supply_aad_prefix(aad.supply_aad_prefix);
if (!aad.aad_prefix.empty()) {
aesGcmV1.__set_aad_prefix(aad.aad_prefix);
}
return aesGcmV1;
}
static inline format::AesGcmCtrV1 ToAesGcmCtrV1Thrift(AadMetadata aad) {
format::AesGcmCtrV1 aesGcmCtrV1;
// aad_file_unique is always set
aesGcmCtrV1.__set_aad_file_unique(aad.aad_file_unique);
aesGcmCtrV1.__set_supply_aad_prefix(aad.supply_aad_prefix);
if (!aad.aad_prefix.empty()) {
aesGcmCtrV1.__set_aad_prefix(aad.aad_prefix);
}
return aesGcmCtrV1;
}
static inline format::EncryptionAlgorithm ToThrift(EncryptionAlgorithm encryption) {
format::EncryptionAlgorithm encryption_algorithm;
if (encryption.algorithm == ParquetCipher::AES_GCM_V1) {
encryption_algorithm.__set_AES_GCM_V1(ToAesGcmV1Thrift(encryption.aad));
} else {
encryption_algorithm.__set_AES_GCM_CTR_V1(ToAesGcmCtrV1Thrift(encryption.aad));
}
return encryption_algorithm;
}
static inline format::SizeStatistics ToThrift(const SizeStatistics& size_stats) {
format::SizeStatistics size_statistics;
size_statistics.__set_definition_level_histogram(size_stats.definition_level_histogram);
size_statistics.__set_repetition_level_histogram(size_stats.repetition_level_histogram);
if (size_stats.unencoded_byte_array_data_bytes.has_value()) {
size_statistics.__set_unencoded_byte_array_data_bytes(
size_stats.unencoded_byte_array_data_bytes.value());
}
return size_statistics;
}
// ----------------------------------------------------------------------
// Thrift struct serialization / deserialization utilities
using ThriftBuffer = apache::thrift::transport::TMemoryBuffer;
class ThriftDeserializer {
public:
explicit ThriftDeserializer(const ReaderProperties& properties)
: ThriftDeserializer(properties.thrift_string_size_limit(),
properties.thrift_container_size_limit()) {}
ThriftDeserializer(int32_t string_size_limit, int32_t container_size_limit)
: string_size_limit_(string_size_limit),
container_size_limit_(container_size_limit) {}
// Deserialize a thrift message from buf/len. buf/len must at least contain
// all the bytes needed to store the thrift message. On return, len will be
// set to the actual length of the header.
template <class T>
void DeserializeMessage(const uint8_t* buf, uint32_t* len, T* deserialized_msg,
Decryptor* decryptor = NULLPTR) {
if (decryptor == NULLPTR) {
// thrift message is not encrypted
DeserializeUnencryptedMessage(buf, len, deserialized_msg);
} else {
// thrift message is encrypted
uint32_t clen;
clen = *len;
if (clen > static_cast<uint32_t>(std::numeric_limits<int32_t>::max())) {
std::stringstream ss;
ss << "Cannot decrypt buffer with length " << clen << ", which overflows int32\n";
throw ParquetException(ss.str());
}
// decrypt
auto decrypted_buffer = AllocateBuffer(
decryptor->pool(), decryptor->PlaintextLength(static_cast<int32_t>(clen)));
std::span<const uint8_t> cipher_buf(buf, clen);
uint32_t decrypted_buffer_len =
decryptor->Decrypt(cipher_buf, decrypted_buffer->mutable_span_as<uint8_t>());
if (decrypted_buffer_len <= 0) {
throw ParquetException("Couldn't decrypt buffer\n");
}
*len = decryptor->CiphertextLength(static_cast<int32_t>(decrypted_buffer_len));
DeserializeUnencryptedMessage(decrypted_buffer->data(), &decrypted_buffer_len,
deserialized_msg);
}
}
private:
// On Thrift 0.14.0+, we want to use TConfiguration to raise the max message size
// limit (ARROW-13655). If we wanted to protect against huge messages, we could
// do it ourselves since we know the message size up front.
std::shared_ptr<ThriftBuffer> CreateReadOnlyMemoryBuffer(uint8_t* buf, uint32_t len) {
#if PARQUET_THRIFT_VERSION_MAJOR > 0 || PARQUET_THRIFT_VERSION_MINOR >= 14
auto conf = std::make_shared<apache::thrift::TConfiguration>();
conf->setMaxMessageSize(std::numeric_limits<int>::max());
return std::make_shared<ThriftBuffer>(buf, len, ThriftBuffer::OBSERVE, conf);
#else
return std::make_shared<ThriftBuffer>(buf, len);
#endif
}
template <class T>
void DeserializeUnencryptedMessage(const uint8_t* buf, uint32_t* len,
T* deserialized_msg) {
// Deserialize msg bytes into c++ thrift msg using memory transport.
auto tmem_transport = CreateReadOnlyMemoryBuffer(const_cast<uint8_t*>(buf), *len);
auto tproto = apache::thrift::protocol::TCompactProtocolT<ThriftBuffer>(
tmem_transport, string_size_limit_, container_size_limit_);
try {
deserialized_msg
->template read<apache::thrift::protocol::TCompactProtocolT<ThriftBuffer>>(
&tproto);
} catch (std::exception& e) {
std::stringstream ss;
ss << "Couldn't deserialize thrift: " << e.what() << "\n";
throw ParquetException(ss.str());
}
uint32_t bytes_left = tmem_transport->available_read();
*len = *len - bytes_left;
}
const int32_t string_size_limit_;
const int32_t container_size_limit_;
};
/// Utility class to serialize thrift objects to a binary format. This object
/// should be reused if possible to reuse the underlying memory.
/// Note: thrift will encode NULLs into the serialized buffer so it is not valid
/// to treat it as a string.
class ThriftSerializer {
public:
explicit ThriftSerializer(int initial_buffer_size = 1024)
: mem_buffer_(new ThriftBuffer(initial_buffer_size)) {
apache::thrift::protocol::TCompactProtocolFactoryT<ThriftBuffer> factory;
protocol_ = factory.getProtocol(mem_buffer_);
}
/// Serialize obj into a memory buffer. The result is returned in buffer/len. The
/// memory returned is owned by this object and will be invalid when another object
/// is serialized.
template <class T>
void SerializeToBuffer(const T* obj, uint32_t* len, uint8_t** buffer) {
SerializeObject(obj);
mem_buffer_->getBuffer(buffer, len);
}
template <class T>
void SerializeToString(const T* obj, std::string* result) {
SerializeObject(obj);
*result = mem_buffer_->getBufferAsString();
}
template <class T>
int64_t Serialize(const T* obj, ArrowOutputStream* out,
Encryptor* encryptor = NULLPTR) {
uint8_t* out_buffer;
uint32_t out_length;
SerializeToBuffer(obj, &out_length, &out_buffer);
// obj is not encrypted
if (encryptor == NULLPTR) {
PARQUET_THROW_NOT_OK(out->Write(out_buffer, out_length));
return static_cast<int64_t>(out_length);
} else { // obj is encrypted
return SerializeEncryptedObj(out, out_buffer, out_length, encryptor);
}
}
private:
template <class T>
void SerializeObject(const T* obj) {
try {
mem_buffer_->resetBuffer();
obj->write(protocol_.get());
} catch (std::exception& e) {
std::stringstream ss;
ss << "Couldn't serialize thrift: " << e.what() << "\n";
throw ParquetException(ss.str());
}
}
int64_t SerializeEncryptedObj(ArrowOutputStream* out, const uint8_t* out_buffer,
uint32_t out_length, Encryptor* encryptor) {
auto cipher_buffer =
AllocateBuffer(encryptor->pool(), encryptor->CiphertextLength(out_length));
std::span<const uint8_t> out_span(out_buffer, out_length);
int32_t cipher_buffer_len =
encryptor->Encrypt(out_span, cipher_buffer->mutable_span_as<uint8_t>());
PARQUET_THROW_NOT_OK(out->Write(cipher_buffer->data(), cipher_buffer_len));
return static_cast<int64_t>(cipher_buffer_len);
}
std::shared_ptr<ThriftBuffer> mem_buffer_;
std::shared_ptr<apache::thrift::protocol::TProtocol> protocol_;
};
} // namespace parquet