@@ -401,6 +401,90 @@ Result<std::shared_ptr<Buffer>> DecompressBuffer(const std::shared_ptr<Buffer>&
401401 return std::move (uncompressed);
402402}
403403
404+ Status DecompressBufferByType (const Buffer& buffer, util::Codec* codec,
405+ std::shared_ptr<Buffer>* out, MemoryPool* pool) {
406+ const uint8_t * data = buffer.data ();
407+ int64_t compressed_size = buffer.size () - sizeof (int64_t );
408+ int64_t uncompressed_size = BitUtil::FromLittleEndian (util::SafeLoadAs<int64_t >(data));
409+
410+ ARROW_ASSIGN_OR_RAISE (auto uncompressed, AllocateBuffer (uncompressed_size, pool));
411+
412+ int64_t actual_decompressed;
413+ ARROW_ASSIGN_OR_RAISE (
414+ actual_decompressed,
415+ codec->Decompress (compressed_size, data + sizeof (int64_t ), uncompressed_size,
416+ uncompressed->mutable_data ()));
417+ if (actual_decompressed != uncompressed_size) {
418+ return Status::Invalid (" Failed to fully decompress buffer, expected "
419+ uncompressed_size, " bytes but decompressed "
420+ actual_decompressed);
421+ }
422+ *out = std::move (uncompressed);
423+ return Status::OK ();
424+ }
425+
426+ Status DecompressBuffersByType (Compression::type compression,
427+ const IpcReadOptions& options,
428+ std::vector<std::shared_ptr<ArrayData>>* arrs,
429+ const std::vector<std::shared_ptr<Field>>& schema_fields) {
430+ ARROW_CHECK_EQ (arrs->size (), schema_fields.size ());
431+
432+ std::unique_ptr<util::Codec> codec;
433+ std::unique_ptr<util::Codec> fastpfor32_codec;
434+ std::unique_ptr<util::Codec> fastpfor64_codec;
435+ ARROW_ASSIGN_OR_RAISE (codec, util::Codec::Create (Compression::LZ4_FRAME));
436+ ARROW_ASSIGN_OR_RAISE (fastpfor32_codec, util::Codec::CreateInt32 (compression));
437+ ARROW_ASSIGN_OR_RAISE (fastpfor64_codec, util::Codec::CreateInt64 (compression));
438+
439+ for (size_t field_idx = 0 ; field_idx < schema_fields.size (); ++field_idx) {
440+ const auto & field = schema_fields[field_idx];
441+ auto & arr = (*arrs)[field_idx];
442+ if (field->type ()->id () == Type::NA) continue ;
443+
444+ const auto & layout_buffers = field->type ()->layout ().buffers ;
445+ for (size_t i = 0 ; i < layout_buffers.size (); ++i) {
446+ const auto & layout = layout_buffers[i];
447+ if (arr->buffers [i] == nullptr ) {
448+ continue ;
449+ }
450+ if (arr->buffers [i]->size () == 0 ) {
451+ continue ;
452+ }
453+ if (arr->buffers [i]->size () < 8 ) {
454+ return Status::Invalid (
455+ " Likely corrupted message, compressed buffers "
456+ " are larger than 8 bytes by construction"
457+ }
458+ auto & buffer = arr->buffers [i];
459+ switch (layout.kind ) {
460+ case DataTypeLayout::BufferKind::FIXED_WIDTH:
461+ if (layout.byte_width == 4 && field->type ()->id () != Type::FLOAT) {
462+ RETURN_NOT_OK (DecompressBufferByType (*buffer, fastpfor32_codec.get (), &buffer,
463+ options.memory_pool ));
464+ } else if (layout.byte_width == 8 && field->type ()->id () != Type::DOUBLE) {
465+ RETURN_NOT_OK (DecompressBufferByType (*buffer, fastpfor64_codec.get (), &buffer,
466+ options.memory_pool ));
467+ } else {
468+ RETURN_NOT_OK (
469+ DecompressBufferByType (*buffer, codec.get (), &buffer, options.memory_pool ));
470+ }
471+ break ;
472+ case DataTypeLayout::BufferKind::BITMAP:
473+ case DataTypeLayout::BufferKind::VARIABLE_WIDTH: {
474+ RETURN_NOT_OK (
475+ DecompressBufferByType (*buffer, codec.get (), &buffer, options.memory_pool ));
476+ break ;
477+ }
478+ case DataTypeLayout::BufferKind::ALWAYS_NULL:
479+ break ;
480+ default :
481+ return Status::Invalid (" Wrong buffer layout."
482+ }
483+ }
484+ }
485+ return arrow::Status::OK ();
486+ }
487+
404488Status DecompressBuffers (Compression::type compression, const IpcReadOptions& options,
405489 ArrayDataVector* fields) {
406490 struct BufferAccumulator {
@@ -482,7 +566,13 @@ Result<std::shared_ptr<RecordBatch>> LoadRecordBatchSubset(
482566 filtered_columns = std::move (columns);
483567 }
484568 if (compression != Compression::UNCOMPRESSED) {
485- RETURN_NOT_OK (DecompressBuffers (compression, options, &filtered_columns));
569+
570+ if (compression == Compression::FASTPFOR) {
571+ RETURN_NOT_OK (
572+ DecompressBuffersByType (compression, options, &filtered_columns, filtered_fields));
573+ } else {
574+ RETURN_NOT_OK (DecompressBuffers (compression, options, &filtered_columns));
575+ }
486576 }
487577
488578 return RecordBatch::Make (filtered_schema, metadata->length (),
0 commit comments