// 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.

#include "arrow/ipc/reader.h"

#include <algorithm>
#include <cstdint>
#include <cstring>
#include <memory>
#include <numeric>
#include <optional>
#include <sstream>
#include <string>
#include <type_traits>
#include <unordered_map>
#include <unordered_set>
#include <utility>
#include <vector>

#include <flatbuffers/flatbuffers.h>  // IWYU pragma: export

#include "arrow/array.h"
#include "arrow/buffer.h"
#include "arrow/extension_type.h"
#include "arrow/io/caching.h"
#include "arrow/io/interfaces.h"
#include "arrow/io/memory.h"
#include "arrow/ipc/message.h"
#include "arrow/ipc/metadata_internal.h"
#include "arrow/ipc/reader_internal.h"
#include "arrow/ipc/writer.h"
#include "arrow/pretty_print.h"
#include "arrow/record_batch.h"
#include "arrow/sparse_tensor.h"
#include "arrow/status.h"
#include "arrow/table.h"
#include "arrow/type.h"
#include "arrow/type_traits.h"
#include "arrow/util/align_util.h"
#include "arrow/util/bit_util.h"
#include "arrow/util/bitmap_ops.h"
#include "arrow/util/checked_cast.h"
#include "arrow/util/compression.h"
#include "arrow/util/endian.h"
#include "arrow/util/fuzz_internal.h"
#include "arrow/util/int_util_overflow.h"
#include "arrow/util/key_value_metadata.h"
#include "arrow/util/logging_internal.h"
#include "arrow/util/parallel.h"
#include "arrow/util/string.h"
#include "arrow/util/thread_pool.h"
#include "arrow/util/ubsan.h"
#include "arrow/util/vector.h"
#include "arrow/visit_type_inline.h"

#include "generated/File_generated.h"  // IWYU pragma: export
#include "generated/Message_generated.h"
#include "generated/Schema_generated.h"
#include "generated/SparseTensor_generated.h"

namespace arrow {

namespace flatbuf = org::apache::arrow::flatbuf;

using internal::AddWithOverflow;
using internal::checked_cast;
using internal::checked_pointer_cast;

namespace ipc {

using internal::FileBlock;
using internal::kArrowMagicBytes;

namespace {

enum class DictionaryKind { New, Delta, Replacement };

Status InvalidMessageType(MessageType expected, MessageType actual) {
  return Status::IOError("Expected IPC message of type ", FormatMessageType(expected),
                         " but got ", FormatMessageType(actual));
}

#define CHECK_MESSAGE_TYPE(expected, actual)           \
  do {                                                 \
    if ((actual) != (expected)) {                      \
      return InvalidMessageType((expected), (actual)); \
    }                                                  \
  } while (0)

#define CHECK_HAS_BODY(message)                                       \
  do {                                                                \
    if ((message).body() == nullptr) {                                \
      return Status::IOError("Expected body in IPC message of type ", \
                             FormatMessageType((message).type()));    \
    }                                                                 \
  } while (0)

#define CHECK_HAS_NO_BODY(message)                                      \
  do {                                                                  \
    if ((message).body_length() != 0) {                                 \
      return Status::IOError("Unexpected body in IPC message of type ", \
                             FormatMessageType((message).type()));      \
    }                                                                   \
  } while (0)

// ----------------------------------------------------------------------
// Record batch read path

/// \brief Structure to keep common arguments to be passed
struct IpcReadContext {
  IpcReadContext(DictionaryMemo* memo, const IpcReadOptions& option, bool swap_endian,
                 MetadataVersion version = MetadataVersion::V5,
                 Compression::type kind = Compression::UNCOMPRESSED)
      : dictionary_memo(memo),
        options(option),
        metadata_version(version),
        compression(kind),
        swap_endian(swap_endian) {}

  DictionaryMemo* dictionary_memo;

  const IpcReadOptions& options;

  MetadataVersion metadata_version;

  Compression::type compression;

  /// \brief LoadRecordBatch() or LoadRecordBatchSubset() swaps endianness of elements
  /// if this flag is true
  const bool swap_endian;
};

/// A collection of ranges to read and pointers to set to those ranges when they are
/// available.  This allows the ArrayLoader to utilize a two pass cache-then-read
/// strategy with a ReadRangeCache
class BatchDataReadRequest {
 public:
  const std::vector<io::ReadRange>& ranges_to_read() const { return ranges_to_read_; }

  void RequestRange(int64_t offset, int64_t length, std::shared_ptr<Buffer>* out) {
    ranges_to_read_.push_back({offset, length});
    destinations_.push_back(out);
  }

  void FulfillRequest(const std::vector<std::shared_ptr<Buffer>>& buffers) {
    for (std::size_t i = 0; i < buffers.size(); i++) {
      *destinations_[i] = buffers[i];
    }
  }

 private:
  std::vector<io::ReadRange> ranges_to_read_;
  std::vector<std::shared_ptr<Buffer>*> destinations_;
};

/// The field_index and buffer_index are incremented based on how much of the
/// batch is "consumed" (through nested data reconstruction, for example)
class ArrayLoader {
 public:
  explicit ArrayLoader(const flatbuf::RecordBatch* metadata,
                       MetadataVersion metadata_version, const IpcReadOptions& options,
                       io::RandomAccessFile* file)
      : metadata_(metadata),
        metadata_version_(metadata_version),
        file_(file),
        file_offset_(0),
        max_recursion_depth_(options.max_recursion_depth) {}

  explicit ArrayLoader(const flatbuf::RecordBatch* metadata,
                       MetadataVersion metadata_version, const IpcReadOptions& options,
                       int64_t file_offset, int64_t file_length)
      : metadata_(metadata),
        metadata_version_(metadata_version),
        file_(nullptr),
        file_offset_(file_offset),
        file_length_(file_length),
        max_recursion_depth_(options.max_recursion_depth) {}

  Status ReadBuffer(int64_t offset, int64_t length, std::shared_ptr<Buffer>* out) {
    // This construct permits overriding GetBuffer at compile time
    if (skip_io_) {
      return Status::OK();
    }
    if (offset < 0) {
      return Status::Invalid("Negative offset for reading buffer ", buffer_index_);
    }
    if (length < 0) {
      return Status::Invalid("Negative length for reading buffer ", buffer_index_);
    }
    auto read_end = AddWithOverflow({offset, length});
    if (!read_end.has_value() || (file_length_.has_value() && read_end > file_length_)) {
      return Status::Invalid("Buffer ", buffer_index_, " exceeds IPC file area");
    }
    if (!bit_util::IsMultipleOf8(offset)) {
      return Status::Invalid("Buffer ", buffer_index_,
                             " did not start on 8-byte aligned offset: ", offset);
    }
    if (file_) {
      return file_->ReadAt(offset, length).Value(out);
    } else {
      if (!AddWithOverflow({read_end.value(), file_offset_}).has_value()) {
        return Status::Invalid("Buffer ", buffer_index_, " exceeds IPC file area");
      }
      read_request_.RequestRange(offset + file_offset_, length, out);
      return Status::OK();
    }
  }

  Status LoadType(const DataType& type) {
    DCHECK_NE(out_, nullptr);
    return VisitTypeInline(type, this);
  }

  Status Load(const Field* field, ArrayData* out) {
    if (max_recursion_depth_ <= 0) {
      return Status::Invalid("Max recursion depth reached");
    }

    field_ = field;
    out_ = out;
    out_->type = field_->type();
    return LoadType(*field_->type());
  }

  Status SkipField(const Field* field) {
    ArrayData dummy;
    skip_io_ = true;
    Status status = Load(field, &dummy);
    skip_io_ = false;
    // GH-37851: Reset state. Load will set `out_` to `&dummy`, which would
    // be a dangling pointer.
    out_ = nullptr;
    return status;
  }

  Status GetBuffer(int buffer_index, std::shared_ptr<Buffer>* out) {
    auto* buffers = metadata_->buffers();
    CHECK_FLATBUFFERS_NOT_NULL(buffers, "RecordBatch.buffers");
    if (buffer_index >= static_cast<int>(buffers->size())) {
      return Status::IOError("buffer_index out of range.");
    }
    const flatbuf::Buffer* buffer = buffers->Get(buffer_index);
    if (buffer->length() == 0) {
      // Should never return a null buffer here.
      // (zero-sized buffer allocations are cheap)
      return AllocateBuffer(0).Value(out);
    } else {
      return ReadBuffer(buffer->offset(), buffer->length(), out);
    }
  }

  Result<int64_t> GetVariadicCount(int i) {
    auto* variadic_counts = metadata_->variadicBufferCounts();
    auto* buffers = metadata_->buffers();
    CHECK_FLATBUFFERS_NOT_NULL(variadic_counts, "RecordBatch.variadicBufferCounts");
    CHECK_FLATBUFFERS_NOT_NULL(buffers, "RecordBatch.buffers");
    if (i >= static_cast<int>(variadic_counts->size())) {
      return Status::IOError("variadic_count_index out of range.");
    }
    int64_t count = variadic_counts->Get(i);
    if (count < 0) {
      return Status::IOError("variadic buffer count must be positive");
    }
    // Detect an excessive variadic buffer count to avoid potential memory blowup
    // (GH-48900).
    const auto max_buffer_count = static_cast<int64_t>(buffers->size()) - buffer_index_;
    if (count > max_buffer_count) {
      return Status::IOError("variadic buffer count exceeds available number of buffers");
    }
    return count;
  }

  Status GetFieldMetadata(int field_index, ArrayData* out) {
    auto nodes = metadata_->nodes();
    CHECK_FLATBUFFERS_NOT_NULL(nodes, "Table.nodes");
    // pop off a field
    if (field_index >= static_cast<int>(nodes->size())) {
      return Status::Invalid("Ran out of field metadata, likely malformed");
    }
    const flatbuf::FieldNode* node = nodes->Get(field_index);

    out->length = node->length();
    out->null_count = node->null_count();
    out->offset = 0;
    return Status::OK();
  }

  Status LoadCommon(Type::type type_id, bool allow_validity_bitmap = true) {
    DCHECK_NE(out_, nullptr);
    // This only contains the length and null count, which we need to figure
    // out what to do with the buffers. For example, if null_count == 0, then
    // we can skip that buffer without reading from shared memory
    RETURN_NOT_OK(GetFieldMetadata(field_index_++, out_));

    if (::arrow::internal::has_variadic_buffers(type_id)) {
      ARROW_ASSIGN_OR_RAISE(auto data_buffer_count,
                            GetVariadicCount(variadic_count_index_++));
      const int64_t start = static_cast<int64_t>(out_->buffers.size());
      // NOTE: this must be done before any other call to `GetBuffer` because
      // BatchDataReadRequest will keep pointers to `std::shared_ptr<Buffer>`
      // objects.
      out_->buffers.resize(start + data_buffer_count);
    }

    if (internal::HasValidityBitmap(type_id, metadata_version_)) {
      // Extract null bitmap which is common to all arrays except for unions
      // and nulls.
      if (out_->null_count != 0) {
        if (allow_validity_bitmap) {
          RETURN_NOT_OK(GetBuffer(buffer_index_, &out_->buffers[0]));
        } else {
          // Caller did not allow this
          return Status::Invalid("Cannot read ", ::arrow::internal::ToTypeName(type_id),
                                 " array with top-level validity bitmap");
        }
      }
      buffer_index_++;
    }

    return Status::OK();
  }

  template <typename TYPE>
  Status LoadPrimitive(Type::type type_id) {
    DCHECK_NE(out_, nullptr);
    out_->buffers.resize(2);

    RETURN_NOT_OK(LoadCommon(type_id));
    if (out_->length > 0) {
      RETURN_NOT_OK(GetBuffer(buffer_index_++, &out_->buffers[1]));
    } else {
      buffer_index_++;
      out_->buffers[1] = std::make_shared<Buffer>(nullptr, 0);
    }
    return Status::OK();
  }

  Status LoadBinary(Type::type type_id) {
    DCHECK_NE(out_, nullptr);
    out_->buffers.resize(3);

    RETURN_NOT_OK(LoadCommon(type_id));
    RETURN_NOT_OK(GetBuffer(buffer_index_++, &out_->buffers[1]));
    return GetBuffer(buffer_index_++, &out_->buffers[2]);
  }

  template <typename TYPE>
  Status LoadList(const TYPE& type) {
    DCHECK_NE(out_, nullptr);
    out_->buffers.resize(2);

    RETURN_NOT_OK(LoadCommon(type.id()));
    RETURN_NOT_OK(GetBuffer(buffer_index_++, &out_->buffers[1]));

    const int num_children = type.num_fields();
    if (num_children != 1) {
      return Status::Invalid("Wrong number of children: ", num_children);
    }

    return LoadChildren(type.fields());
  }

  template <typename TYPE>
  Status LoadListView(const TYPE& type) {
    out_->buffers.resize(3);

    RETURN_NOT_OK(LoadCommon(type.id()));
    RETURN_NOT_OK(GetBuffer(buffer_index_++, &out_->buffers[1]));
    RETURN_NOT_OK(GetBuffer(buffer_index_++, &out_->buffers[2]));

    const int num_children = type.num_fields();
    if (num_children != 1) {
      return Status::Invalid("Wrong number of children: ", num_children);
    }

    return LoadChildren(type.fields());
  }

  Status LoadChildren(const std::vector<std::shared_ptr<Field>>& child_fields) {
    DCHECK_NE(out_, nullptr);
    ArrayData* parent = out_;

    parent->child_data.resize(child_fields.size());
    for (int i = 0; i < static_cast<int>(child_fields.size()); ++i) {
      parent->child_data[i] = std::make_shared<ArrayData>();
      --max_recursion_depth_;
      RETURN_NOT_OK(Load(child_fields[i].get(), parent->child_data[i].get()));
      ++max_recursion_depth_;
    }
    out_ = parent;
    return Status::OK();
  }

  Status Visit(const NullType& type) {
    out_->buffers.resize(1);

    // ARROW-6379: NullType has no buffers in the IPC payload
    return GetFieldMetadata(field_index_++, out_);
  }

  template <typename T>
  enable_if_t<std::is_base_of<FixedWidthType, T>::value &&
                  !std::is_base_of<FixedSizeBinaryType, T>::value &&
                  !std::is_base_of<DictionaryType, T>::value,
              Status>
  Visit(const T& type) {
    return LoadPrimitive<T>(type.id());
  }

  template <typename T>
  enable_if_base_binary<T, Status> Visit(const T& type) {
    return LoadBinary(type.id());
  }

  Status Visit(const BinaryViewType& type) {
    out_->buffers.resize(2);

    RETURN_NOT_OK(LoadCommon(type.id()));  // also initializes variadic buffers
    for (int64_t i = 1; i < static_cast<int64_t>(out_->buffers.size()); ++i) {
      RETURN_NOT_OK(GetBuffer(buffer_index_++, &out_->buffers[i]));
    }
    return Status::OK();
  }

  Status Visit(const FixedSizeBinaryType& type) {
    out_->buffers.resize(2);
    RETURN_NOT_OK(LoadCommon(type.id()));
    return GetBuffer(buffer_index_++, &out_->buffers[1]);
  }

  template <typename T>
  enable_if_var_size_list<T, Status> Visit(const T& type) {
    return LoadList(type);
  }

  template <typename T>
  enable_if_list_view<T, Status> Visit(const T& type) {
    return LoadListView(type);
  }

  Status Visit(const MapType& type) {
    RETURN_NOT_OK(LoadList(type));
    return MapArray::ValidateChildData(out_->child_data);
  }

  Status Visit(const FixedSizeListType& type) {
    out_->buffers.resize(1);

    RETURN_NOT_OK(LoadCommon(type.id()));

    const int num_children = type.num_fields();
    if (num_children != 1) {
      return Status::Invalid("Wrong number of children: ", num_children);
    }

    return LoadChildren(type.fields());
  }

  Status Visit(const StructType& type) {
    out_->buffers.resize(1);
    RETURN_NOT_OK(LoadCommon(type.id()));
    return LoadChildren(type.fields());
  }

  Status Visit(const UnionType& type) {
    int n_buffers = type.mode() == UnionMode::SPARSE ? 2 : 3;
    out_->buffers.resize(n_buffers);

    // With metadata V4, we can get a validity bitmap. The bitmap may be there
    // if we're loading eagerly, or it might be scheduled for loading if we're
    // using a BatchDataReadRequest.
    //
    // Trying to fix up union data to do without the top-level validity bitmap
    // is hairy:
    // - type ids must be rewritten to all have valid values (even for former
    //   null slots)
    // - sparse union children must have their validity bitmaps rewritten
    //   by ANDing the top-level validity bitmap
    // - dense union children must be rewritten (at least one of them)
    //   to insert the required null slots that were formerly omitted
    //
    // So instead we disallow validity bitmaps.
    RETURN_NOT_OK(LoadCommon(type.id(), /*allow_validity_bitmap=*/false));
    out_->null_count = 0;

    if (out_->length > 0) {
      RETURN_NOT_OK(GetBuffer(buffer_index_, &out_->buffers[1]));
      if (type.mode() == UnionMode::DENSE) {
        RETURN_NOT_OK(GetBuffer(buffer_index_ + 1, &out_->buffers[2]));
      }
    }
    buffer_index_ += n_buffers - 1;
    return LoadChildren(type.fields());
  }

  Status Visit(const DictionaryType& type) {
    // out_->dictionary will be filled later in ResolveDictionaries()
    return LoadType(*type.index_type());
  }

  Status Visit(const RunEndEncodedType& type) {
    out_->buffers.resize(1);
    RETURN_NOT_OK(LoadCommon(type.id()));
    return LoadChildren(type.fields());
  }

  Status Visit(const ExtensionType& type) { return LoadType(*type.storage_type()); }

  BatchDataReadRequest& read_request() { return read_request_; }

 private:
  const flatbuf::RecordBatch* metadata_;
  const MetadataVersion metadata_version_;
  io::RandomAccessFile* file_;
  int64_t file_offset_;
  std::optional<int64_t> file_length_;
  int max_recursion_depth_;
  int buffer_index_ = 0;
  int field_index_ = 0;
  bool skip_io_ = false;
  int variadic_count_index_ = 0;

  BatchDataReadRequest read_request_;
  const Field* field_ = nullptr;
  ArrayData* out_ = nullptr;
};

Result<std::shared_ptr<Buffer>> DecompressBuffer(const std::shared_ptr<Buffer>& buf,
                                                 const IpcReadOptions& options,
                                                 util::Codec* codec) {
  if (buf == nullptr || buf->size() == 0) {
    return buf;
  }

  if (buf->size() < 8) {
    return Status::Invalid(
        "Likely corrupted message, compressed buffers "
        "are larger than 8 bytes by construction");
  }

  const uint8_t* data = buf->data();
  int64_t compressed_size = buf->size() - sizeof(int64_t);
  int64_t uncompressed_size = bit_util::FromLittleEndian(util::SafeLoadAs<int64_t>(data));

  if (uncompressed_size == -1) {
    return SliceBuffer(buf, sizeof(int64_t), compressed_size);
  }

  ARROW_ASSIGN_OR_RAISE(auto uncompressed,
                        AllocateBuffer(uncompressed_size, options.memory_pool));

  ARROW_ASSIGN_OR_RAISE(
      int64_t actual_decompressed,
      codec->Decompress(compressed_size, data + sizeof(int64_t), uncompressed_size,
                        uncompressed->mutable_data()));
  if (actual_decompressed != uncompressed_size) {
    return Status::Invalid("Failed to fully decompress buffer, expected ",
                           uncompressed_size, " bytes but decompressed ",
                           actual_decompressed);
  }

  // R build with openSUSE155 requires an explicit shared_ptr construction
  return std::shared_ptr<Buffer>(std::move(uncompressed));
}

Status DecompressBuffers(Compression::type compression, const IpcReadOptions& options,
                         ArrayDataVector* fields) {
  struct BufferAccumulator {
    using BufferPtrVector = std::vector<std::shared_ptr<Buffer>*>;

    void AppendFrom(const ArrayDataVector& fields) {
      for (const auto& field : fields) {
        for (auto& buffer : field->buffers) {
          buffers_.push_back(&buffer);
        }
        AppendFrom(field->child_data);
      }
      // Dictionary buffers are decompressed separately (see ReadDictionary).
    }

    BufferPtrVector Get(const ArrayDataVector& fields) && {
      AppendFrom(fields);
      return std::move(buffers_);
    }

    BufferPtrVector buffers_;
  };

  // Flatten all buffers
  auto buffers = BufferAccumulator{}.Get(*fields);

  std::unique_ptr<util::Codec> codec;
  ARROW_ASSIGN_OR_RAISE(codec, util::Codec::Create(compression));

  return ::arrow::internal::OptionalParallelFor(
      options.use_threads, static_cast<int>(buffers.size()), [&](int i) {
        ARROW_ASSIGN_OR_RAISE(*buffers[i],
                              DecompressBuffer(*buffers[i], options, codec.get()));
        return Status::OK();
      });
}

// Helper class to run post-ArrayLoader steps:
// buffer decompression, dictionary resolution, buffer re-alignment.
struct RecordBatchLoader {
  Result<std::shared_ptr<RecordBatch>> CreateRecordBatch(ArrayDataVector columns) {
    ARROW_ASSIGN_OR_RAISE(auto filtered_columns, CreateColumns(std::move(columns)));

    std::shared_ptr<Schema> filtered_schema;
    if (!inclusion_mask_.empty()) {
      FieldVector filtered_fields;
      for (int i = 0; i < schema_->num_fields(); ++i) {
        if (inclusion_mask_[i]) {
          filtered_fields.push_back(schema_->field(i));
        }
      }
      filtered_schema = schema(std::move(filtered_fields), schema_->metadata());
    } else {
      filtered_schema = schema_;
    }

    return RecordBatch::Make(std::move(filtered_schema), batch_length_,
                             std::move(filtered_columns));
  }

  Result<ArrayDataVector> CreateColumns(ArrayDataVector columns,
                                        bool resolve_dictionaries = true) {
    if (resolve_dictionaries) {
      // Dictionary resolution needs to happen on the unfiltered columns,
      // because fields are mapped structurally (by path in the original schema).
      RETURN_NOT_OK(ResolveDictionaries(columns, *context_.dictionary_memo,
                                        context_.options.memory_pool));
    }

    ArrayDataVector filtered_columns;
    if (!inclusion_mask_.empty()) {
      FieldVector filtered_fields;
      for (int i = 0; i < schema_->num_fields(); ++i) {
        if (inclusion_mask_[i]) {
          DCHECK_NE(columns[i], nullptr);
          filtered_columns.push_back(std::move(columns[i]));
        }
      }
      columns.clear();
    } else {
      filtered_columns = std::move(columns);
    }

    if (context_.compression != Compression::UNCOMPRESSED) {
      RETURN_NOT_OK(
          DecompressBuffers(context_.compression, context_.options, &filtered_columns));
    }

    // Swap endian if necessary
    if (context_.swap_endian) {
      for (auto& column : filtered_columns) {
        ARROW_ASSIGN_OR_RAISE(
            column, arrow::internal::SwapEndianArrayData(std::move(column),
                                                         context_.options.memory_pool));
      }
    }
    if (context_.options.ensure_alignment != Alignment::kAnyAlignment) {
      for (auto& column : filtered_columns) {
        ARROW_ASSIGN_OR_RAISE(
            column,
            util::EnsureAlignment(
                std::move(column),
                // The numerical value of the enum is taken literally as byte alignment
                static_cast<int64_t>(context_.options.ensure_alignment),
                context_.options.memory_pool));
      }
    }
    return filtered_columns;
  }

  IpcReadContext context_;
  std::shared_ptr<Schema> schema_;
  int64_t batch_length_;
  std::vector<bool> inclusion_mask_;
};

Result<std::shared_ptr<RecordBatch>> LoadRecordBatchSubset(
    const flatbuf::RecordBatch* metadata, const std::shared_ptr<Schema>& schema,
    const std::vector<bool>* inclusion_mask, const IpcReadContext& context,
    io::RandomAccessFile* file) {
  ArrayLoader loader(metadata, context.metadata_version, context.options, file);
  ArrayDataVector columns(schema->num_fields());

  for (int i = 0; i < schema->num_fields(); ++i) {
    const Field& field = *schema->field(i);
    if (!inclusion_mask || (*inclusion_mask)[i]) {
      // Read field
      auto column = std::make_shared<ArrayData>();
      RETURN_NOT_OK(loader.Load(&field, column.get()));
      if (metadata->length() != column->length) {
        return Status::IOError("Array length did not match record batch length");
      }
      columns[i] = std::move(column);
    } else {
      // Skip field. This logic must be executed to advance the state of the
      // loader to the next field
      RETURN_NOT_OK(loader.SkipField(&field));
    }
  }

  RecordBatchLoader batch_loader{context, schema, metadata->length(),
                                 inclusion_mask ? *inclusion_mask : std::vector<bool>{}};
  return batch_loader.CreateRecordBatch(std::move(columns));
}

Result<std::shared_ptr<RecordBatch>> LoadRecordBatch(
    const flatbuf::RecordBatch* metadata, const std::shared_ptr<Schema>& schema,
    const std::vector<bool>& inclusion_mask, const IpcReadContext& context,
    io::RandomAccessFile* file) {
  if (inclusion_mask.empty()) {
    return LoadRecordBatchSubset(metadata, schema, /*inclusion_mask=*/nullptr, context,
                                 file);
  } else {
    return LoadRecordBatchSubset(metadata, schema, &inclusion_mask, context, file);
  }
}

// ----------------------------------------------------------------------
// Array loading

Status GetCompression(const flatbuf::RecordBatch* batch, Compression::type* out) {
  *out = Compression::UNCOMPRESSED;
  const flatbuf::BodyCompression* compression = batch->compression();
  if (compression != nullptr) {
    if (compression->method() !=
        flatbuf::BodyCompressionMethod::BodyCompressionMethod_BUFFER) {
      // Forward compatibility
      return Status::Invalid("This library only supports BUFFER compression method");
    }

    if (compression->codec() == flatbuf::CompressionType::CompressionType_LZ4_FRAME) {
      *out = Compression::LZ4_FRAME;
    } else if (compression->codec() == flatbuf::CompressionType::CompressionType_ZSTD) {
      *out = Compression::ZSTD;
    } else {
      return Status::Invalid("Unsupported codec in RecordBatch::compression metadata");
    }
    return Status::OK();
  }
  return Status::OK();
}

Status GetCompressionExperimental(const flatbuf::Message* message,
                                  Compression::type* out) {
  *out = Compression::UNCOMPRESSED;
  if (message->custom_metadata() != nullptr) {
    // TODO: Ensure this deserialization only ever happens once
    std::shared_ptr<KeyValueMetadata> metadata;
    RETURN_NOT_OK(internal::GetKeyValueMetadata(message->custom_metadata(), &metadata));
    int index = metadata->FindKey("ARROW:experimental_compression");
    if (index != -1) {
      // Arrow 0.17 stored string in upper case, internal utils now require lower case
      auto name = arrow::internal::AsciiToLower(metadata->value(index));
      ARROW_ASSIGN_OR_RAISE(*out, util::Codec::GetCompressionType(name));
    }
    return internal::CheckCompressionSupported(*out);
  }
  return Status::OK();
}

Status ReadContiguousPayload(io::InputStream* file, std::unique_ptr<Message>* message) {
  ARROW_ASSIGN_OR_RAISE(*message, ReadMessage(file));
  if (*message == nullptr) {
    return Status::Invalid("Unable to read metadata at offset");
  }
  return Status::OK();
}

Result<RecordBatchWithMetadata> ReadRecordBatchInternal(
    const Buffer& metadata, const std::shared_ptr<Schema>& schema,
    const std::vector<bool>& inclusion_mask, IpcReadContext& context,
    io::RandomAccessFile* file) {
  const flatbuf::Message* message = nullptr;
  RETURN_NOT_OK(internal::VerifyMessage(metadata.data(), metadata.size(), &message));
  auto batch = message->header_as_RecordBatch();
  if (batch == nullptr) {
    return Status::IOError(
        "Header-type of flatbuffer-encoded Message is not RecordBatch.");
  }

  Compression::type compression;
  RETURN_NOT_OK(GetCompression(batch, &compression));
  if (context.compression == Compression::UNCOMPRESSED &&
      message->version() == flatbuf::MetadataVersion::MetadataVersion_V4) {
    // Possibly obtain codec information from experimental serialization format
    // in 0.17.x
    RETURN_NOT_OK(GetCompressionExperimental(message, &compression));
  }
  context.compression = compression;
  context.metadata_version = internal::GetMetadataVersion(message->version());

  std::shared_ptr<KeyValueMetadata> custom_metadata;
  if (message->custom_metadata() != nullptr) {
    RETURN_NOT_OK(
        internal::GetKeyValueMetadata(message->custom_metadata(), &custom_metadata));
  }
  ARROW_ASSIGN_OR_RAISE(auto record_batch,
                        LoadRecordBatch(batch, schema, inclusion_mask, context, file));
  return RecordBatchWithMetadata{record_batch, custom_metadata};
}

// If we are selecting only certain fields, populate an inclusion mask for fast lookups.
// Additionally, drop deselected fields from the reader's schema.
Status GetInclusionMaskAndOutSchema(const std::shared_ptr<Schema>& full_schema,
                                    const std::vector<int>& included_indices,
                                    std::vector<bool>* inclusion_mask,
                                    std::shared_ptr<Schema>* out_schema) {
  inclusion_mask->clear();
  if (included_indices.empty()) {
    *out_schema = full_schema;
    return Status::OK();
  }

  inclusion_mask->resize(full_schema->num_fields(), false);

  auto included_indices_sorted = included_indices;
  std::sort(included_indices_sorted.begin(), included_indices_sorted.end());

  FieldVector included_fields;
  for (int i : included_indices_sorted) {
    // Ignore out of bounds indices
    if (i < 0 || i >= full_schema->num_fields()) {
      return Status::Invalid("Out of bounds field index: ", i);
    }

    if (inclusion_mask->at(i)) continue;

    inclusion_mask->at(i) = true;
    included_fields.push_back(full_schema->field(i));
  }

  *out_schema = schema(std::move(included_fields), full_schema->endianness(),
                       full_schema->metadata());
  return Status::OK();
}

Status UnpackSchemaMessage(const void* opaque_schema, const IpcReadOptions& options,
                           DictionaryMemo* dictionary_memo,
                           std::shared_ptr<Schema>* schema,
                           std::shared_ptr<Schema>* out_schema,
                           std::vector<bool>* field_inclusion_mask, bool* swap_endian) {
  RETURN_NOT_OK(internal::GetSchema(opaque_schema, dictionary_memo, schema));

  // If we are selecting only certain fields, populate the inclusion mask now
  // for fast lookups
  RETURN_NOT_OK(GetInclusionMaskAndOutSchema(*schema, options.included_fields,
                                             field_inclusion_mask, out_schema));
  *swap_endian = options.ensure_native_endian && !out_schema->get()->is_native_endian();
  if (*swap_endian) {
    // create a new schema with native endianness before swapping endian in ArrayData
    *schema = schema->get()->WithEndianness(Endianness::Native);
    *out_schema = out_schema->get()->WithEndianness(Endianness::Native);
  }
  return Status::OK();
}

Status UnpackSchemaMessage(const Message& message, const IpcReadOptions& options,
                           DictionaryMemo* dictionary_memo,
                           std::shared_ptr<Schema>* schema,
                           std::shared_ptr<Schema>* out_schema,
                           std::vector<bool>* field_inclusion_mask, bool* swap_endian) {
  CHECK_MESSAGE_TYPE(MessageType::SCHEMA, message.type());
  CHECK_HAS_NO_BODY(message);

  return UnpackSchemaMessage(message.header(), options, dictionary_memo, schema,
                             out_schema, field_inclusion_mask, swap_endian);
}

Status ReadDictionary(const Buffer& metadata, IpcReadContext context,
                      DictionaryKind* kind, io::RandomAccessFile* file) {
  const flatbuf::Message* message = nullptr;
  RETURN_NOT_OK(internal::VerifyMessage(metadata.data(), metadata.size(), &message));
  const auto dictionary_batch = message->header_as_DictionaryBatch();
  if (dictionary_batch == nullptr) {
    return Status::IOError(
        "Header-type of flatbuffer-encoded Message is not DictionaryBatch.");
  }

  // The dictionary is embedded in a record batch with a single column
  const auto batch_meta = dictionary_batch->data();

  CHECK_FLATBUFFERS_NOT_NULL(batch_meta, "DictionaryBatch.data");

  RETURN_NOT_OK(GetCompression(batch_meta, &context.compression));
  if (context.compression == Compression::UNCOMPRESSED &&
      message->version() == flatbuf::MetadataVersion::MetadataVersion_V4) {
    // Possibly obtain codec information from experimental serialization format
    // in 0.17.x
    RETURN_NOT_OK(GetCompressionExperimental(message, &context.compression));
  }

  const int64_t id = dictionary_batch->id();

  // Look up the dictionary value type, which must have been added to the
  // DictionaryMemo already prior to invoking this function
  ARROW_ASSIGN_OR_RAISE(auto value_type, context.dictionary_memo->GetDictionaryType(id));

  // Load the dictionary data from the dictionary batch
  ArrayLoader loader(batch_meta, internal::GetMetadataVersion(message->version()),
                     context.options, file);
  auto dict_data = std::make_shared<ArrayData>();
  const Field dummy_field("", value_type);
  RETURN_NOT_OK(loader.Load(&dummy_field, dict_data.get()));

  // Run post-load steps: buffer decompression, etc.
  RecordBatchLoader batch_loader{context, /*schema=*/nullptr, batch_meta->length(),
                                 /*inclusion_mask=*/std::vector<bool>{}};
  ARROW_ASSIGN_OR_RAISE(
      auto dict_columns,
      batch_loader.CreateColumns({dict_data}, /*resolve_dictionaries=*/false));
  DCHECK_EQ(dict_columns.size(), 1);
  dict_data = dict_columns[0];

  if (dictionary_batch->isDelta()) {
    if (kind != nullptr) {
      *kind = DictionaryKind::Delta;
    }
    return context.dictionary_memo->AddDictionaryDelta(id, dict_data);
  }
  ARROW_ASSIGN_OR_RAISE(bool inserted,
                        context.dictionary_memo->AddOrReplaceDictionary(id, dict_data));
  if (kind != nullptr) {
    *kind = inserted ? DictionaryKind::New : DictionaryKind::Replacement;
  }
  return Status::OK();
}

Status ReadDictionary(const Message& message, const IpcReadContext& context,
                      DictionaryKind* kind) {
  // Only invoke this method if we already know we have a dictionary message
  DCHECK_EQ(message.type(), MessageType::DICTIONARY_BATCH);
  CHECK_HAS_BODY(message);
  ARROW_ASSIGN_OR_RAISE(auto reader, Buffer::GetReader(message.body()));
  return ReadDictionary(*message.metadata(), context, kind, reader.get());
}

}  // namespace

Result<std::shared_ptr<RecordBatch>> ReadRecordBatch(
    const Buffer& metadata, const std::shared_ptr<Schema>& schema,
    const DictionaryMemo* dictionary_memo, const IpcReadOptions& options,
    io::RandomAccessFile* file) {
  std::shared_ptr<Schema> out_schema;
  // Empty means do not use
  std::vector<bool> inclusion_mask;
  IpcReadContext context(const_cast<DictionaryMemo*>(dictionary_memo), options, false);
  RETURN_NOT_OK(GetInclusionMaskAndOutSchema(schema, context.options.included_fields,
                                             &inclusion_mask, &out_schema));
  ARROW_ASSIGN_OR_RAISE(
      auto batch_and_custom_metadata,
      ReadRecordBatchInternal(metadata, schema, inclusion_mask, context, file));
  return batch_and_custom_metadata.batch;
}

Result<std::shared_ptr<RecordBatch>> ReadRecordBatch(
    const std::shared_ptr<Schema>& schema, const DictionaryMemo* dictionary_memo,
    const IpcReadOptions& options, io::InputStream* file) {
  std::unique_ptr<Message> message;
  RETURN_NOT_OK(ReadContiguousPayload(file, &message));
  CHECK_HAS_BODY(*message);
  ARROW_ASSIGN_OR_RAISE(auto reader, Buffer::GetReader(message->body()));
  return ReadRecordBatch(*message->metadata(), schema, dictionary_memo, options,
                         reader.get());
}

Result<std::shared_ptr<RecordBatch>> ReadRecordBatch(
    const Message& message, const std::shared_ptr<Schema>& schema,
    const DictionaryMemo* dictionary_memo, const IpcReadOptions& options) {
  CHECK_MESSAGE_TYPE(MessageType::RECORD_BATCH, message.type());
  CHECK_HAS_BODY(message);
  ARROW_ASSIGN_OR_RAISE(auto reader, Buffer::GetReader(message.body()));
  return ReadRecordBatch(*message.metadata(), schema, dictionary_memo, options,
                         reader.get());
}

// Streaming format decoder
class StreamDecoderInternal : public MessageDecoderListener {
 public:
  enum State {
    SCHEMA,
    INITIAL_DICTIONARIES,
    RECORD_BATCHES,
    EOS,
  };

  explicit StreamDecoderInternal(std::shared_ptr<Listener> listener,
                                 IpcReadOptions options)
      : listener_(std::move(listener)),
        options_(std::move(options)),
        state_(State::SCHEMA),
        field_inclusion_mask_(),
        num_required_initial_dictionaries_(0),
        num_read_initial_dictionaries_(0),
        dictionary_memo_(),
        schema_(nullptr),
        filtered_schema_(nullptr),
        stats_(),
        swap_endian_(false) {}

  Status OnMessageDecoded(std::unique_ptr<Message> message) override {
    ++stats_.num_messages;
    switch (state_) {
      case State::SCHEMA:
        ARROW_RETURN_NOT_OK(OnSchemaMessageDecoded(std::move(message)));
        break;
      case State::INITIAL_DICTIONARIES:
        ARROW_RETURN_NOT_OK(OnInitialDictionaryMessageDecoded(std::move(message)));
        break;
      case State::RECORD_BATCHES:
        ARROW_RETURN_NOT_OK(OnRecordBatchMessageDecoded(std::move(message)));
        break;
      case State::EOS:
        break;
    }
    return Status::OK();
  }

  Status OnEOS() override {
    state_ = State::EOS;
    return listener_->OnEOS();
  }

  std::shared_ptr<Listener> listener() const { return listener_; }

  Listener* raw_listener() const { return listener_.get(); }

  IpcReadOptions options() const { return options_; }

  State state() const { return state_; }

  std::shared_ptr<Schema> schema() const { return filtered_schema_; }

  ReadStats stats() const { return stats_; }

  int num_required_initial_dictionaries() const {
    return num_required_initial_dictionaries_;
  }

  int num_read_initial_dictionaries() const { return num_read_initial_dictionaries_; }

 private:
  Status OnSchemaMessageDecoded(std::unique_ptr<Message> message) {
    RETURN_NOT_OK(UnpackSchemaMessage(*message, options_, &dictionary_memo_, &schema_,
                                      &filtered_schema_, &field_inclusion_mask_,
                                      &swap_endian_));

    num_required_initial_dictionaries_ = dictionary_memo_.fields().num_dicts();
    num_read_initial_dictionaries_ = 0;
    if (num_required_initial_dictionaries_ == 0) {
      state_ = State::RECORD_BATCHES;
      RETURN_NOT_OK(listener_->OnSchemaDecoded(schema_, filtered_schema_));
    } else {
      state_ = State::INITIAL_DICTIONARIES;
    }
    return Status::OK();
  }

  Status OnInitialDictionaryMessageDecoded(std::unique_ptr<Message> message) {
    if (message->type() != MessageType::DICTIONARY_BATCH) {
      return Status::Invalid("IPC stream did not have the expected number (",
                             num_required_initial_dictionaries_,
                             ") of dictionaries at the start of the stream");
    }
    RETURN_NOT_OK(ReadDictionary(*message));
    num_read_initial_dictionaries_++;
    if (num_read_initial_dictionaries_ == num_required_initial_dictionaries_) {
      state_ = State::RECORD_BATCHES;
      ARROW_RETURN_NOT_OK(listener_->OnSchemaDecoded(schema_, filtered_schema_));
    }
    return Status::OK();
  }

  Status OnRecordBatchMessageDecoded(std::unique_ptr<Message> message) {
    if (message->type() == MessageType::DICTIONARY_BATCH) {
      return ReadDictionary(*message);
    } else {
      CHECK_HAS_BODY(*message);
      ARROW_ASSIGN_OR_RAISE(auto reader, Buffer::GetReader(message->body()));
      IpcReadContext context(&dictionary_memo_, options_, swap_endian_);
      ARROW_ASSIGN_OR_RAISE(
          auto batch_with_metadata,
          ReadRecordBatchInternal(*message->metadata(), schema_, field_inclusion_mask_,
                                  context, reader.get()));
      ++stats_.num_record_batches;
      return listener_->OnRecordBatchWithMetadataDecoded(batch_with_metadata);
    }
  }

  // Read dictionary from dictionary batch
  Status ReadDictionary(const Message& message) {
    DictionaryKind kind;
    IpcReadContext context(&dictionary_memo_, options_, swap_endian_);
    RETURN_NOT_OK(::arrow::ipc::ReadDictionary(message, context, &kind));
    ++stats_.num_dictionary_batches;
    switch (kind) {
      case DictionaryKind::New:
        break;
      case DictionaryKind::Delta:
        ++stats_.num_dictionary_deltas;
        break;
      case DictionaryKind::Replacement:
        ++stats_.num_replaced_dictionaries;
        break;
    }
    return Status::OK();
  }

  std::shared_ptr<Listener> listener_;
  const IpcReadOptions options_;
  State state_;
  std::vector<bool> field_inclusion_mask_;
  int num_required_initial_dictionaries_;
  int num_read_initial_dictionaries_;
  DictionaryMemo dictionary_memo_;
  std::shared_ptr<Schema> schema_;
  std::shared_ptr<Schema> filtered_schema_;
  ReadStats stats_;
  bool swap_endian_;
};

// ----------------------------------------------------------------------
// RecordBatchStreamReader implementation

class RecordBatchStreamReaderImpl : public RecordBatchStreamReader,
                                    public StreamDecoderInternal {
 public:
  RecordBatchStreamReaderImpl(std::unique_ptr<MessageReader> message_reader,
                              const IpcReadOptions& options)
      : RecordBatchStreamReader(),
        StreamDecoderInternal(std::make_shared<CollectListener>(), options),
        message_reader_(std::move(message_reader)) {}

  Status Init() {
    // Read schema
    ARROW_ASSIGN_OR_RAISE(auto message, message_reader_->ReadNextMessage());
    if (!message) {
      return Status::Invalid("Tried reading schema message, was null or length 0");
    }
    return OnMessageDecoded(std::move(message));
  }

  Status ReadNext(std::shared_ptr<RecordBatch>* batch) override {
    ARROW_ASSIGN_OR_RAISE(auto batch_with_metadata, ReadNext());
    *batch = std::move(batch_with_metadata.batch);
    return Status::OK();
  }

  Result<RecordBatchWithMetadata> ReadNext() override {
    auto collect_listener = checked_cast<CollectListener*>(raw_listener());
    while (collect_listener->num_record_batches() == 0 &&
           state() != StreamDecoderInternal::State::EOS) {
      ARROW_ASSIGN_OR_RAISE(auto message, message_reader_->ReadNextMessage());
      if (!message) {  // End of stream
        if (state() == StreamDecoderInternal::State::INITIAL_DICTIONARIES) {
          if (num_read_initial_dictionaries() == 0) {
            // ARROW-6006: If we fail to find any dictionaries in the
            // stream, then it may be that the stream has a schema
            // but no actual data. In such case we communicate that
            // we were unable to find the dictionaries (but there was
            // no failure otherwise), so the caller can decide what
            // to do
            return RecordBatchWithMetadata{nullptr, nullptr};
          } else {
            // ARROW-6126, the stream terminated before receiving the
            // expected number of dictionaries
            return Status::Invalid(
                "IPC stream ended without reading the "
                "expected number (",
                num_required_initial_dictionaries(), ") of dictionaries");
          }
        } else {
          return RecordBatchWithMetadata{nullptr, nullptr};
        }
      }
      ARROW_RETURN_NOT_OK(OnMessageDecoded(std::move(message)));
    }
    return collect_listener->PopRecordBatchWithMetadata();
  }

  std::shared_ptr<Schema> schema() const override {
    return StreamDecoderInternal::schema();
  }

  ReadStats stats() const override { return StreamDecoderInternal::stats(); }

 private:
  std::unique_ptr<MessageReader> message_reader_;
};

// ----------------------------------------------------------------------
// Stream reader constructors

Result<std::shared_ptr<RecordBatchStreamReader>> RecordBatchStreamReader::Open(
    std::unique_ptr<MessageReader> message_reader, const IpcReadOptions& options) {
  // Private ctor
  auto result =
      std::make_shared<RecordBatchStreamReaderImpl>(std::move(message_reader), options);
  RETURN_NOT_OK(result->Init());
  return result;
}

Result<std::shared_ptr<RecordBatchStreamReader>> RecordBatchStreamReader::Open(
    io::InputStream* stream, const IpcReadOptions& options) {
  return Open(MessageReader::Open(stream), options);
}

Result<std::shared_ptr<RecordBatchStreamReader>> RecordBatchStreamReader::Open(
    const std::shared_ptr<io::InputStream>& stream, const IpcReadOptions& options) {
  return Open(MessageReader::Open(stream), options);
}

// ----------------------------------------------------------------------
// Reader implementation

namespace {

// Common functions used in both the random-access file reader and the
// asynchronous generator
Result<FileBlock> FileBlockFromFlatbuffer(const flatbuf::Block* fb_block,
                                          int64_t max_offset) {
  auto block =
      FileBlock{fb_block->offset(), fb_block->metaDataLength(), fb_block->bodyLength()};
  if (block.metadata_length < 0 || block.body_length < 0 || block.offset < 0) {
    return Status::IOError("Invalid Block in IPC file footer");
  }
  auto block_end =
      AddWithOverflow<int64_t>({block.offset, block.metadata_length, block.body_length});
  if (!block_end.has_value() || block_end > max_offset) {
    return Status::IOError("Invalid Block in IPC file footer");
  }
  return block;
}

Status CheckAligned(const FileBlock& block) {
  if (!bit_util::IsMultipleOf8(block.offset) ||
      !bit_util::IsMultipleOf8(block.metadata_length) ||
      !bit_util::IsMultipleOf8(block.body_length)) {
    return Status::Invalid("Unaligned block in IPC file");
  }
  return Status::OK();
}

template <typename MessagePtr>
Result<MessagePtr> CheckBodyLength(MessagePtr message, const FileBlock& block) {
  if (message->body_length() != block.body_length) {
    return Status::Invalid(
        "Mismatching body length for IPC message "
        "(Block.bodyLength: ",
        block.body_length, " vs. Message.bodyLength: ", message->body_length(), ")");
  }
  // NOTE: we cannot check metadata length as easily as we would have to account
  // for the additional IPC signalisation (such as optional continuation bytes).
  return message;
}

Result<std::unique_ptr<Message>> ReadMessageFromBlock(
    const FileBlock& block, io::RandomAccessFile* file,
    const FieldsLoaderFunction& fields_loader) {
  RETURN_NOT_OK(CheckAligned(block));
  if (fields_loader) {
    ARROW_ASSIGN_OR_RAISE(auto message, ReadMessage(block.offset, block.metadata_length,
                                                    file, fields_loader));
    return CheckBodyLength(std::move(message), block);
  } else {
    ARROW_ASSIGN_OR_RAISE(auto message, ReadMessage(block.offset, block.metadata_length,
                                                    block.body_length, file));
    return CheckBodyLength(std::move(message), block);
  }
}

Future<std::shared_ptr<Message>> ReadMessageFromBlockAsync(
    const FileBlock& block, io::RandomAccessFile* file, const io::IOContext& io_context) {
  RETURN_NOT_OK(CheckAligned(block));
  return ReadMessageAsync(block.offset, block.metadata_length, block.body_length, file,
                          io_context)
      .Then([block](std::shared_ptr<Message> message) {
        return CheckBodyLength(std::move(message), block);
      });
}

class RecordBatchFileReaderImpl;

/// A generator of record batches.
///
/// All batches are yielded in order.
class WholeIpcFileRecordBatchGenerator {
 public:
  using Item = std::shared_ptr<RecordBatch>;

  explicit WholeIpcFileRecordBatchGenerator(
      std::shared_ptr<RecordBatchFileReaderImpl> state,
      std::shared_ptr<io::internal::ReadRangeCache> cached_source,
      const io::IOContext& io_context, arrow::internal::Executor* executor)
      : state_(std::move(state)),
        cached_source_(std::move(cached_source)),
        io_context_(io_context),
        executor_(executor),
        index_(0) {}

  Future<Item> operator()();
  Future<std::shared_ptr<Message>> ReadBlock(const FileBlock& block);

  static Status ReadDictionaries(
      RecordBatchFileReaderImpl* state,
      std::vector<std::shared_ptr<Message>> dictionary_messages);
  static Result<std::shared_ptr<RecordBatch>> ReadRecordBatch(
      RecordBatchFileReaderImpl* state, Message* message);

 private:
  std::shared_ptr<RecordBatchFileReaderImpl> state_;
  std::shared_ptr<io::internal::ReadRangeCache> cached_source_;
  io::IOContext io_context_;
  arrow::internal::Executor* executor_;
  int index_;
  // Odd Future type, but this lets us use All() easily
  Future<> read_dictionaries_;
};

/// A generator of record batches for use when reading
/// a subset of columns from the file.
///
/// All batches are yielded in order.
class SelectiveIpcFileRecordBatchGenerator {
 public:
  using Item = std::shared_ptr<RecordBatch>;

  explicit SelectiveIpcFileRecordBatchGenerator(
      std::shared_ptr<RecordBatchFileReaderImpl> state)
      : state_(std::move(state)), index_(0) {}

  Future<Item> operator()();

 private:
  std::shared_ptr<RecordBatchFileReaderImpl> state_;
  int index_;
};

class RecordBatchFileReaderImpl : public RecordBatchFileReader {
 public:
  RecordBatchFileReaderImpl() : file_(NULLPTR), footer_offset_(0), footer_(NULLPTR) {}

  int num_record_batches() const override {
    return static_cast<int>(internal::FlatBuffersVectorSize(footer_->recordBatches()));
  }

  MetadataVersion version() const override {
    return internal::GetMetadataVersion(footer_->version());
  }

  static Status LoadFieldsSubset(const flatbuf::RecordBatch* metadata,
                                 const IpcReadOptions& options,
                                 io::RandomAccessFile* file,
                                 const std::shared_ptr<Schema>& schema,
                                 const std::vector<bool>* inclusion_mask,
                                 MetadataVersion metadata_version = MetadataVersion::V5) {
    ArrayLoader loader(metadata, metadata_version, options, file);
    for (int i = 0; i < schema->num_fields(); ++i) {
      const Field& field = *schema->field(i);
      if (!inclusion_mask || (*inclusion_mask)[i]) {
        // Read field
        ArrayData column;
        RETURN_NOT_OK(loader.Load(&field, &column));
        if (metadata->length() != column.length) {
          return Status::IOError("Array length did not match record batch length");
        }
      } else {
        // Skip field. This logic must be executed to advance the state of the
        // loader to the next field
        RETURN_NOT_OK(loader.SkipField(&field));
      }
    }
    return Status::OK();
  }

  Future<std::shared_ptr<RecordBatch>> ReadRecordBatchAsync(int i) {
    DCHECK_GE(i, 0);
    DCHECK_LT(i, num_record_batches());

    auto cached_metadata = cached_metadata_.find(i);
    if (cached_metadata != cached_metadata_.end()) {
      return ReadCachedRecordBatch(i, cached_metadata->second);
    }

    return Status::Invalid(
        "Asynchronous record batch reading is only supported after a call to "
        "PreBufferMetadata or PreBufferBatches");
  }

  Result<std::shared_ptr<RecordBatch>> ReadRecordBatch(int i) override {
    ARROW_ASSIGN_OR_RAISE(auto batch_with_metadata, ReadRecordBatchWithCustomMetadata(i));
    return batch_with_metadata.batch;
  }

  Result<RecordBatchWithMetadata> ReadRecordBatchWithCustomMetadata(int i) override {
    DCHECK_GE(i, 0);
    DCHECK_LT(i, num_record_batches());

    auto cached_metadata = cached_metadata_.find(i);
    if (cached_metadata != cached_metadata_.end()) {
      auto result = ReadCachedRecordBatch(i, cached_metadata->second).result();
      ARROW_ASSIGN_OR_RAISE(auto batch, result);
      ARROW_ASSIGN_OR_RAISE(auto message_obj, cached_metadata->second.result());
      ARROW_ASSIGN_OR_RAISE(auto message, GetFlatbufMessage(message_obj));
      std::shared_ptr<KeyValueMetadata> custom_metadata;
      if (message->custom_metadata() != nullptr) {
        RETURN_NOT_OK(
            internal::GetKeyValueMetadata(message->custom_metadata(), &custom_metadata));
      }
      return RecordBatchWithMetadata{std::move(batch), std::move(custom_metadata)};
    }

    RETURN_NOT_OK(WaitForDictionaryReadFinished());

    FieldsLoaderFunction fields_loader = {};
    if (!field_inclusion_mask_.empty()) {
      auto& schema = schema_;
      auto& inclusion_mask = field_inclusion_mask_;
      auto& read_options = options_;
      fields_loader = [schema, inclusion_mask, read_options](const void* metadata,
                                                             io::RandomAccessFile* file) {
        return LoadFieldsSubset(static_cast<const flatbuf::RecordBatch*>(metadata),
                                read_options, file, schema, &inclusion_mask);
      };
    }
    ARROW_ASSIGN_OR_RAISE(auto block, GetRecordBatchBlock(i));
    ARROW_ASSIGN_OR_RAISE(auto message, ReadMessageFromBlock(block, fields_loader));

    CHECK_HAS_BODY(*message);
    ARROW_ASSIGN_OR_RAISE(auto reader, Buffer::GetReader(message->body()));
    IpcReadContext context(&dictionary_memo_, options_, swap_endian_);
    ARROW_ASSIGN_OR_RAISE(
        auto batch_with_metadata,
        ReadRecordBatchInternal(*message->metadata(), schema_, field_inclusion_mask_,
                                context, reader.get()));
    stats_.num_record_batches.fetch_add(1, std::memory_order_relaxed);
    return batch_with_metadata;
  }

  Result<int64_t> CountRows() override {
    int64_t total = 0;
    for (int i = 0; i < num_record_batches(); i++) {
      ARROW_ASSIGN_OR_RAISE(auto block, GetRecordBatchBlock(i));
      ARROW_ASSIGN_OR_RAISE(auto outer_message, ReadMessageFromBlock(block));
      auto metadata = outer_message->metadata();
      const flatbuf::Message* message = nullptr;
      RETURN_NOT_OK(
          internal::VerifyMessage(metadata->data(), metadata->size(), &message));
      auto batch = message->header_as_RecordBatch();
      if (batch == nullptr) {
        return Status::IOError(
            "Header-type of flatbuffer-encoded Message is not RecordBatch.");
      }
      total += batch->length();
    }
    return total;
  }

  Status Open(const std::shared_ptr<io::RandomAccessFile>& file, int64_t footer_offset,
              const IpcReadOptions& options) {
    owned_file_ = file;
    metadata_cache_ = std::make_shared<io::internal::ReadRangeCache>(
        file, file->io_context(), options.pre_buffer_cache_options);
    return Open(file.get(), footer_offset, options);
  }

  Status Open(io::RandomAccessFile* file, int64_t footer_offset,
              const IpcReadOptions& options) {
    // The metadata_cache_ may have already been constructed with an owned file in the
    // owning overload of Open
    if (!metadata_cache_) {
      metadata_cache_ = std::make_shared<io::internal::ReadRangeCache>(
          file, file->io_context(), options.pre_buffer_cache_options);
    }
    file_ = file;
    options_ = options;
    footer_offset_ = footer_offset;
    RETURN_NOT_OK(ReadFooter());

    // Get the schema and record any observed dictionaries
    RETURN_NOT_OK(UnpackSchemaMessage(footer_->schema(), options, &dictionary_memo_,
                                      &schema_, &out_schema_, &field_inclusion_mask_,
                                      &swap_endian_));
    stats_.num_messages.fetch_add(1, std::memory_order_relaxed);
    return Status::OK();
  }

  Future<> OpenAsync(const std::shared_ptr<io::RandomAccessFile>& file,
                     int64_t footer_offset, const IpcReadOptions& options) {
    owned_file_ = file;
    metadata_cache_ = std::make_shared<io::internal::ReadRangeCache>(
        file, file->io_context(), options.pre_buffer_cache_options);
    return OpenAsync(file.get(), footer_offset, options);
  }

  Future<> OpenAsync(io::RandomAccessFile* file, int64_t footer_offset,
                     const IpcReadOptions& options) {
    // The metadata_cache_ may have already been constructed with an owned file in the
    // owning overload of OpenAsync
    if (!metadata_cache_) {
      metadata_cache_ = std::make_shared<io::internal::ReadRangeCache>(
          file, file->io_context(), options.pre_buffer_cache_options);
    }
    file_ = file;
    options_ = options;
    footer_offset_ = footer_offset;
    auto cpu_executor = ::arrow::internal::GetCpuThreadPool();
    auto self = std::dynamic_pointer_cast<RecordBatchFileReaderImpl>(shared_from_this());
    return ReadFooterAsync(cpu_executor).Then([self, options]() -> Status {
      // Get the schema and record any observed dictionaries
      RETURN_NOT_OK(UnpackSchemaMessage(
          self->footer_->schema(), options, &self->dictionary_memo_, &self->schema_,
          &self->out_schema_, &self->field_inclusion_mask_, &self->swap_endian_));
      self->stats_.num_messages.fetch_add(1, std::memory_order_relaxed);
      return Status::OK();
    });
  }

  std::shared_ptr<Schema> schema() const override { return out_schema_; }

  std::shared_ptr<const KeyValueMetadata> metadata() const override { return metadata_; }

  ReadStats stats() const override { return stats_.poll(); }

  Result<AsyncGenerator<std::shared_ptr<RecordBatch>>> GetRecordBatchGenerator(
      const bool coalesce, const io::IOContext& io_context,
      const io::CacheOptions cache_options,
      arrow::internal::Executor* executor) override {
    auto state = std::dynamic_pointer_cast<RecordBatchFileReaderImpl>(shared_from_this());
    // Prebuffering causes us to use a lot of futures which, at the moment,
    // can only slow things down when we are doing zero-copy in-memory reads.
    //
    // Prebuffering's read patterns are also slightly worse than the alternative
    // when doing whole-file reads because the logic is not in place to recognize
    // we can just read the entire file up-front
    if (!options_.included_fields.empty() &&
        options_.included_fields.size() != schema_->fields().size() &&
        !file_->supports_zero_copy()) {
      RETURN_NOT_OK(state->PreBufferMetadata({}));
      return SelectiveIpcFileRecordBatchGenerator(std::move(state));
    }

    std::shared_ptr<io::internal::ReadRangeCache> cached_source;
    if (coalesce && !file_->supports_zero_copy()) {
      if (!owned_file_) return Status::Invalid("Cannot coalesce without an owned file");
      // Since the user is asking for all fields then we can cache the entire
      // file (up to the footer)
      cached_source = std::make_shared<io::internal::ReadRangeCache>(file_, io_context,
                                                                     cache_options);
      RETURN_NOT_OK(cached_source->Cache({{0, footer_offset_}}));
    }
    return WholeIpcFileRecordBatchGenerator(std::move(state), std::move(cached_source),
                                            io_context, executor);
  }

  Status DoPreBufferMetadata(const std::vector<int>& indices) {
    RETURN_NOT_OK(CacheMetadata(indices));
    RETURN_NOT_OK(EnsureDictionaryReadStarted());
    Future<> all_metadata_ready = WaitForMetadatas(indices);
    for (int index : indices) {
      Future<std::shared_ptr<Message>> metadata_loaded =
          all_metadata_ready.Then([this, index]() -> Result<std::shared_ptr<Message>> {
            stats_.num_messages.fetch_add(1, std::memory_order_relaxed);
            ARROW_ASSIGN_OR_RAISE(FileBlock block, GetRecordBatchBlock(index));
            ARROW_ASSIGN_OR_RAISE(
                std::shared_ptr<Buffer> metadata,
                metadata_cache_->Read({block.offset, block.metadata_length}));
            return ReadMessage(std::move(metadata), nullptr);
          });
      cached_metadata_.emplace(index, metadata_loaded);
    }
    return Status::OK();
  }

  std::vector<int> AllIndices() const {
    std::vector<int> all_indices(num_record_batches());
    std::iota(all_indices.begin(), all_indices.end(), 0);
    return all_indices;
  }

  Status PreBufferMetadata(const std::vector<int>& indices) override {
    if (indices.size() == 0) {
      return DoPreBufferMetadata(AllIndices());
    } else {
      return DoPreBufferMetadata(indices);
    }
  }

 private:
  friend class WholeIpcFileRecordBatchGenerator;

  struct AtomicReadStats {
    std::atomic<int64_t> num_messages{0};
    std::atomic<int64_t> num_record_batches{0};
    std::atomic<int64_t> num_dictionary_batches{0};
    std::atomic<int64_t> num_dictionary_deltas{0};
    std::atomic<int64_t> num_replaced_dictionaries{0};

    /// \brief Capture a copy of the current counters
    ReadStats poll() const {
      ReadStats stats;
      stats.num_messages = num_messages.load(std::memory_order_relaxed);
      stats.num_record_batches = num_record_batches.load(std::memory_order_relaxed);
      stats.num_dictionary_batches =
          num_dictionary_batches.load(std::memory_order_relaxed);
      stats.num_dictionary_deltas = num_dictionary_deltas.load(std::memory_order_relaxed);
      stats.num_replaced_dictionaries =
          num_replaced_dictionaries.load(std::memory_order_relaxed);
      return stats;
    }
  };

  Result<FileBlock> GetRecordBatchBlock(int i) const {
    return FileBlockFromFlatbuffer(footer_->recordBatches()->Get(i), footer_offset_);
  }

  Result<FileBlock> GetDictionaryBlock(int i) const {
    return FileBlockFromFlatbuffer(footer_->dictionaries()->Get(i), footer_offset_);
  }

  Result<std::unique_ptr<Message>> ReadMessageFromBlock(
      const FileBlock& block, const FieldsLoaderFunction& fields_loader = {}) {
    ARROW_ASSIGN_OR_RAISE(auto message,
                          arrow::ipc::ReadMessageFromBlock(block, file_, fields_loader));
    stats_.num_messages.fetch_add(1, std::memory_order_relaxed);
    return message;
  }

  Status ReadDictionaries() {
    // Read all the dictionaries
    std::vector<std::shared_ptr<Message>> messages(num_dictionaries());
    for (int i = 0; i < num_dictionaries(); ++i) {
      ARROW_ASSIGN_OR_RAISE(FileBlock block, GetDictionaryBlock(i));
      ARROW_ASSIGN_OR_RAISE(messages[i], ReadMessageFromBlock(block));
    }
    return ReadDictionaries(messages);
  }

  Status ReadDictionaries(
      const std::vector<std::shared_ptr<Message>>& dictionary_messages) {
    DCHECK_EQ(dictionary_messages.size(), static_cast<size_t>(num_dictionaries()));
    IpcReadContext context(&dictionary_memo_, options_, swap_endian_);
    for (int i = 0; i < num_dictionaries(); ++i) {
      RETURN_NOT_OK(ReadOneDictionary(i, dictionary_messages[i].get(), context));
    }
    return Status::OK();
  }

  Status ReadOneDictionary(int dict_index, Message* message,
                           const IpcReadContext& context) {
    CHECK_HAS_BODY(*message);
    ARROW_ASSIGN_OR_RAISE(auto reader, Buffer::GetReader(message->body()));
    DictionaryKind kind;
    RETURN_NOT_OK(ReadDictionary(*message->metadata(), context, &kind, reader.get()));
    if (kind == DictionaryKind::Replacement) {
      return Status::Invalid("Unsupported dictionary replacement in IPC file");
    } else if (kind == DictionaryKind::Delta) {
      stats_.num_dictionary_deltas.fetch_add(1, std::memory_order_relaxed);
    }
    stats_.num_dictionary_batches.fetch_add(1, std::memory_order_relaxed);
    return Status::OK();
  }

  Status AddDictionaryRanges(std::vector<io::ReadRange>* ranges) const {
    // Adds all dictionaries to the range cache
    for (int i = 0; i < num_dictionaries(); ++i) {
      ARROW_ASSIGN_OR_RAISE(FileBlock block, GetDictionaryBlock(i));
      ranges->push_back({block.offset, block.metadata_length + block.body_length});
    }
    return Status::OK();
  }

  Status AddMetadataRanges(const std::vector<int>& indices,
                           std::vector<io::ReadRange>* ranges) {
    for (int index : indices) {
      ARROW_ASSIGN_OR_RAISE(FileBlock block, GetRecordBatchBlock(index));
      ranges->push_back({block.offset, block.metadata_length});
    }
    return Status::OK();
  }

  Status CacheMetadata(const std::vector<int>& indices) {
    std::vector<io::ReadRange> ranges;
    if (!read_dictionaries_) {
      RETURN_NOT_OK(AddDictionaryRanges(&ranges));
    }
    RETURN_NOT_OK(AddMetadataRanges(indices, &ranges));
    return metadata_cache_->Cache(std::move(ranges));
  }

  Status EnsureDictionaryReadStarted() {
    if (!dictionary_load_finished_.is_valid()) {
      read_dictionaries_ = true;
      std::vector<io::ReadRange> ranges;
      RETURN_NOT_OK(AddDictionaryRanges(&ranges));
      dictionary_load_finished_ =
          metadata_cache_->WaitFor(std::move(ranges)).Then([this] {
            return ReadDictionaries();
          });
    }
    return Status::OK();
  }

  Status WaitForDictionaryReadFinished() {
    if (!read_dictionaries_) {
      RETURN_NOT_OK(ReadDictionaries());
      read_dictionaries_ = true;
      return Status::OK();
    }
    if (dictionary_load_finished_.is_valid()) {
      return dictionary_load_finished_.status();
    }
    // Dictionaries were previously loaded synchronously
    return Status::OK();
  }

  Future<> WaitForMetadatas(const std::vector<int>& indices) {
    std::vector<io::ReadRange> ranges;
    RETURN_NOT_OK(AddMetadataRanges(indices, &ranges));
    return metadata_cache_->WaitFor(std::move(ranges));
  }

  Result<IpcReadContext> GetIpcReadContext(const flatbuf::Message* message,
                                           const flatbuf::RecordBatch* batch) {
    IpcReadContext context(&dictionary_memo_, options_, swap_endian_);
    Compression::type compression;
    RETURN_NOT_OK(GetCompression(batch, &compression));
    if (context.compression == Compression::UNCOMPRESSED &&
        message->version() == flatbuf::MetadataVersion::MetadataVersion_V4) {
      // Possibly obtain codec information from experimental serialization format
      // in 0.17.x
      RETURN_NOT_OK(GetCompressionExperimental(message, &compression));
    }
    context.compression = compression;
    context.metadata_version = internal::GetMetadataVersion(message->version());
    return context;
  }

  Result<const flatbuf::RecordBatch*> GetBatchFromMessage(
      const flatbuf::Message* message) {
    auto batch = message->header_as_RecordBatch();
    if (batch == nullptr) {
      return Status::IOError(
          "Header-type of flatbuffer-encoded Message is not RecordBatch.");
    }
    return batch;
  }

  Result<const flatbuf::Message*> GetFlatbufMessage(
      const std::shared_ptr<Message>& message) {
    const Buffer& metadata = *message->metadata();
    const flatbuf::Message* flatbuf_message = nullptr;
    RETURN_NOT_OK(
        internal::VerifyMessage(metadata.data(), metadata.size(), &flatbuf_message));
    return flatbuf_message;
  }

  struct CachedRecordBatchReadContext {
    CachedRecordBatchReadContext(std::shared_ptr<Schema> sch,
                                 const flatbuf::RecordBatch* batch,
                                 IpcReadContext context, io::RandomAccessFile* file,
                                 std::shared_ptr<io::RandomAccessFile> owned_file,
                                 int64_t block_data_offset, int64_t block_data_length)
        : schema(std::move(sch)),
          context(std::move(context)),
          file(file),
          owned_file(std::move(owned_file)),
          loader(batch, context.metadata_version, context.options, block_data_offset,
                 block_data_length),
          columns(schema->num_fields()),
          cache(file, file->io_context(), io::CacheOptions::LazyDefaults()),
          length(batch->length()) {}

    Status CalculateLoadRequest() {
      std::shared_ptr<Schema> out_schema;
      RETURN_NOT_OK(GetInclusionMaskAndOutSchema(schema, context.options.included_fields,
                                                 &inclusion_mask, &out_schema));
      for (int i = 0; i < schema->num_fields(); ++i) {
        const Field& field = *schema->field(i);
        if (inclusion_mask.empty() || inclusion_mask[i]) {
          // Read field
          auto column = std::make_shared<ArrayData>();
          RETURN_NOT_OK(loader.Load(&field, column.get()));
          if (length != column->length) {
            return Status::IOError("Array length did not match record batch length");
          }
          columns[i] = std::move(column);
        } else {
          // Skip field. This logic must be executed to advance the state of the
          // loader to the next field
          RETURN_NOT_OK(loader.SkipField(&field));
        }
      }
      return Status::OK();
    }

    Future<> ReadAsync() {
      RETURN_NOT_OK(cache.Cache(loader.read_request().ranges_to_read()));
      return cache.WaitFor(loader.read_request().ranges_to_read());
    }

    Result<std::shared_ptr<RecordBatch>> CreateRecordBatch() {
      std::vector<std::shared_ptr<Buffer>> buffers;
      for (const auto& range_to_read : loader.read_request().ranges_to_read()) {
        ARROW_ASSIGN_OR_RAISE(auto buffer, cache.Read(range_to_read));
        buffers.push_back(std::move(buffer));
      }
      loader.read_request().FulfillRequest(buffers);

      RecordBatchLoader batch_loader{context, schema, length, std::move(inclusion_mask)};
      return batch_loader.CreateRecordBatch(std::move(columns));
    }

    std::shared_ptr<Schema> schema;
    IpcReadContext context;
    io::RandomAccessFile* file;
    std::shared_ptr<io::RandomAccessFile> owned_file;

    ArrayLoader loader;
    ArrayDataVector columns;
    io::internal::ReadRangeCache cache;
    int64_t length;
    std::vector<bool> inclusion_mask;
  };

  Future<std::shared_ptr<RecordBatch>> ReadCachedRecordBatch(
      int index, Future<std::shared_ptr<Message>> message_fut) {
    stats_.num_record_batches.fetch_add(1, std::memory_order_relaxed);
    return dictionary_load_finished_.Then([message_fut] { return message_fut; })
        .Then([this, index](const std::shared_ptr<Message>& message_obj)
                  -> Future<std::shared_ptr<RecordBatch>> {
          ARROW_ASSIGN_OR_RAISE(auto block, GetRecordBatchBlock(index));
          ARROW_ASSIGN_OR_RAISE(auto message, GetFlatbufMessage(message_obj));
          ARROW_ASSIGN_OR_RAISE(auto batch, GetBatchFromMessage(message));
          ARROW_ASSIGN_OR_RAISE(auto context, GetIpcReadContext(message, batch));

          auto read_context = std::make_shared<CachedRecordBatchReadContext>(
              schema_, batch, std::move(context), file_, owned_file_,
              block.offset + static_cast<int64_t>(block.metadata_length),
              block.body_length);
          RETURN_NOT_OK(read_context->CalculateLoadRequest());
          return read_context->ReadAsync().Then(
              [read_context] { return read_context->CreateRecordBatch(); });
        });
  }

  Status ReadFooter() {
    auto fut = ReadFooterAsync(/*executor=*/nullptr);
    return fut.status();
  }

  Future<> ReadFooterAsync(arrow::internal::Executor* executor) {
    constexpr int32_t kMagicSize = static_cast<int>(kArrowMagicBytes.size());

    if (footer_offset_ <= kMagicSize * 2 + 4) {
      return Status::Invalid("File is too small: ", footer_offset_);
    }

    int file_end_size = static_cast<int>(kMagicSize + sizeof(int32_t));
    auto self = std::dynamic_pointer_cast<RecordBatchFileReaderImpl>(shared_from_this());
    auto read_magic = file_->ReadAsync(footer_offset_ - file_end_size, file_end_size);
    if (executor) read_magic = executor->Transfer(std::move(read_magic));
    return read_magic
        .Then([=](const std::shared_ptr<Buffer>& buffer)
                  -> Future<std::shared_ptr<Buffer>> {
          const int64_t expected_footer_size = kMagicSize + sizeof(int32_t);
          if (buffer->size() < expected_footer_size) {
            return Status::Invalid("Unable to read ", expected_footer_size,
                                   "from end of file");
          }

          const auto magic_start = buffer->data() + sizeof(int32_t);
          if (std::string_view(reinterpret_cast<const char*>(magic_start), kMagicSize) !=
              kArrowMagicBytes) {
            return Status::Invalid("Not an Arrow file");
          }

          int32_t footer_length = bit_util::FromLittleEndian(
              *reinterpret_cast<const int32_t*>(buffer->data()));

          if (footer_length <= 0 ||
              footer_length > self->footer_offset_ - kMagicSize * 2 - 4) {
            return Status::Invalid("File is smaller than indicated metadata size");
          }

          // Now read the footer
          auto read_footer = self->file_->ReadAsync(
              self->footer_offset_ - footer_length - file_end_size, footer_length);
          if (executor) read_footer = executor->Transfer(std::move(read_footer));
          return read_footer;
        })
        .Then([=](const std::shared_ptr<Buffer>& buffer) -> Status {
          self->footer_buffer_ = buffer;
          const auto data = self->footer_buffer_->data();
          const auto size = self->footer_buffer_->size();
          if (!internal::VerifyFlatbuffers<flatbuf::Footer>(data, size)) {
            return Status::IOError("Verification of flatbuffer-encoded Footer failed.");
          }
          self->footer_ = flatbuf::GetFooter(data);

          auto fb_metadata = self->footer_->custom_metadata();
          if (fb_metadata != nullptr) {
            std::shared_ptr<KeyValueMetadata> md;
            RETURN_NOT_OK(internal::GetKeyValueMetadata(fb_metadata, &md));
            self->metadata_ = std::move(md);  // const-ify
          }
          return Status::OK();
        });
  }

  int num_dictionaries() const {
    return static_cast<int>(internal::FlatBuffersVectorSize(footer_->dictionaries()));
  }

  io::RandomAccessFile* file_;
  IpcReadOptions options_;
  std::vector<bool> field_inclusion_mask_;

  std::shared_ptr<io::RandomAccessFile> owned_file_;

  // The location where the Arrow file layout ends. May be the end of the file
  // or some other location if embedded in a larger file.
  int64_t footer_offset_;

  // Footer metadata
  std::shared_ptr<Buffer> footer_buffer_;
  const flatbuf::Footer* footer_;
  std::shared_ptr<const KeyValueMetadata> metadata_;

  bool read_dictionaries_ = false;
  DictionaryMemo dictionary_memo_;

  // Reconstructed schema, including any read dictionaries
  std::shared_ptr<Schema> schema_;
  // Schema with deselected fields dropped
  std::shared_ptr<Schema> out_schema_;

  AtomicReadStats stats_;
  std::shared_ptr<io::internal::ReadRangeCache> metadata_cache_;
  std::unordered_set<int> cached_data_blocks_;
  Future<> dictionary_load_finished_;
  std::unordered_map<int, Future<std::shared_ptr<Message>>> cached_metadata_;
  std::unordered_map<int, Future<>> cached_data_requests_;

  bool swap_endian_;
};

Future<SelectiveIpcFileRecordBatchGenerator::Item>
SelectiveIpcFileRecordBatchGenerator::operator()() {
  int index = index_++;
  if (index >= state_->num_record_batches()) {
    return IterationEnd<SelectiveIpcFileRecordBatchGenerator::Item>();
  }
  return state_->ReadRecordBatchAsync(index);
}

Future<WholeIpcFileRecordBatchGenerator::Item>
WholeIpcFileRecordBatchGenerator::operator()() {
  auto state = state_;
  if (!read_dictionaries_.is_valid()) {
    if (state->dictionary_load_finished_.is_valid()) {
      // PreBufferMetadata has started reading dictionaries in the background
      read_dictionaries_ = state->dictionary_load_finished_;
    } else {
      // Start reading dictionaries
      std::vector<Future<std::shared_ptr<Message>>> messages(state->num_dictionaries());
      for (int i = 0; i < state->num_dictionaries(); i++) {
        ARROW_ASSIGN_OR_RAISE(auto block, state->GetDictionaryBlock(i));
        messages[i] = ReadBlock(block);
      }
      auto read_messages = All(std::move(messages));
      if (executor_) read_messages = executor_->Transfer(read_messages);
      read_dictionaries_ = read_messages.Then(
          [=](const std::vector<Result<std::shared_ptr<Message>>>& maybe_messages)
              -> Status {
            ARROW_ASSIGN_OR_RAISE(auto messages,
                                  arrow::internal::UnwrapOrRaise(maybe_messages));
            return state->ReadDictionaries(messages);
          });
    }
  }
  if (index_ >= state_->num_record_batches()) {
    return Future<Item>::MakeFinished(IterationTraits<Item>::End());
  }
  ARROW_ASSIGN_OR_RAISE(auto block, state->GetRecordBatchBlock(index_++));
  auto read_message = ReadBlock(block);
  auto read_messages = read_dictionaries_.Then([read_message]() { return read_message; });
  // Force transfer. This may be wasteful in some cases, but ensures we get off the
  // I/O threads as soon as possible, and ensures we don't decode record batches
  // synchronously in the case that the message read has already finished.
  if (executor_) {
    auto executor = executor_;
    return read_messages.Then(
        [=](const std::shared_ptr<Message>& message) -> Future<Item> {
          return DeferNotOk(executor->Submit(
              [=]() { return ReadRecordBatch(state.get(), message.get()); }));
        });
  }
  return read_messages.Then([=](const std::shared_ptr<Message>& message) -> Result<Item> {
    return ReadRecordBatch(state.get(), message.get());
  });
}

Future<std::shared_ptr<Message>> WholeIpcFileRecordBatchGenerator::ReadBlock(
    const FileBlock& block) {
  if (cached_source_) {
    auto cached_source = cached_source_;
    io::ReadRange range{block.offset, block.metadata_length + block.body_length};
    auto pool = state_->options_.memory_pool;
    return cached_source->WaitFor({range}).Then(
        [cached_source, pool, range]() -> Result<std::shared_ptr<Message>> {
          ARROW_ASSIGN_OR_RAISE(auto buffer, cached_source->Read(range));
          io::BufferReader stream(std::move(buffer));
          return ReadMessage(&stream, pool);
        });
  } else {
    return ReadMessageFromBlockAsync(block, state_->file_, io_context_);
  }
}

Result<std::shared_ptr<RecordBatch>> WholeIpcFileRecordBatchGenerator::ReadRecordBatch(
    RecordBatchFileReaderImpl* state, Message* message) {
  CHECK_HAS_BODY(*message);
  ARROW_ASSIGN_OR_RAISE(auto reader, Buffer::GetReader(message->body()));
  IpcReadContext context(&state->dictionary_memo_, state->options_, state->swap_endian_);
  ARROW_ASSIGN_OR_RAISE(
      auto batch_with_metadata,
      ReadRecordBatchInternal(*message->metadata(), state->schema_,
                              state->field_inclusion_mask_, context, reader.get()));
  return batch_with_metadata.batch;
}

}  // namespace

Result<std::shared_ptr<RecordBatchFileReader>> RecordBatchFileReader::Open(
    io::RandomAccessFile* file, const IpcReadOptions& options) {
  ARROW_ASSIGN_OR_RAISE(int64_t footer_offset, file->GetSize());
  return Open(file, footer_offset, options);
}

Result<std::shared_ptr<RecordBatchFileReader>> RecordBatchFileReader::Open(
    io::RandomAccessFile* file, int64_t footer_offset, const IpcReadOptions& options) {
  auto result = std::make_shared<RecordBatchFileReaderImpl>();
  RETURN_NOT_OK(result->Open(file, footer_offset, options));
  return result;
}

Result<std::shared_ptr<RecordBatchFileReader>> RecordBatchFileReader::Open(
    const std::shared_ptr<io::RandomAccessFile>& file, const IpcReadOptions& options) {
  ARROW_ASSIGN_OR_RAISE(int64_t footer_offset, file->GetSize());
  return Open(file, footer_offset, options);
}

Result<std::shared_ptr<RecordBatchFileReader>> RecordBatchFileReader::Open(
    const std::shared_ptr<io::RandomAccessFile>& file, int64_t footer_offset,
    const IpcReadOptions& options) {
  auto result = std::make_shared<RecordBatchFileReaderImpl>();
  RETURN_NOT_OK(result->Open(file, footer_offset, options));
  return result;
}

Future<std::shared_ptr<RecordBatchFileReader>> RecordBatchFileReader::OpenAsync(
    const std::shared_ptr<io::RandomAccessFile>& file, const IpcReadOptions& options) {
  ARROW_ASSIGN_OR_RAISE(int64_t footer_offset, file->GetSize());
  return OpenAsync(file, footer_offset, options);
}

Future<std::shared_ptr<RecordBatchFileReader>> RecordBatchFileReader::OpenAsync(
    io::RandomAccessFile* file, const IpcReadOptions& options) {
  ARROW_ASSIGN_OR_RAISE(int64_t footer_offset, file->GetSize());
  return OpenAsync(file, footer_offset, options);
}

Future<std::shared_ptr<RecordBatchFileReader>> RecordBatchFileReader::OpenAsync(
    const std::shared_ptr<io::RandomAccessFile>& file, int64_t footer_offset,
    const IpcReadOptions& options) {
  auto result = std::make_shared<RecordBatchFileReaderImpl>();
  return result->OpenAsync(file, footer_offset, options)
      .Then([=]() -> Result<std::shared_ptr<RecordBatchFileReader>> { return result; });
}

Future<std::shared_ptr<RecordBatchFileReader>> RecordBatchFileReader::OpenAsync(
    io::RandomAccessFile* file, int64_t footer_offset, const IpcReadOptions& options) {
  auto result = std::make_shared<RecordBatchFileReaderImpl>();
  return result->OpenAsync(file, footer_offset, options)
      .Then([=]() -> Result<std::shared_ptr<RecordBatchFileReader>> { return result; });
}

Result<RecordBatchVector> RecordBatchFileReader::ToRecordBatches() {
  RecordBatchVector batches;
  const auto n = num_record_batches();
  for (int i = 0; i < n; ++i) {
    ARROW_ASSIGN_OR_RAISE(auto batch, ReadRecordBatch(i));
    batches.emplace_back(std::move(batch));
  }
  return batches;
}

Result<std::shared_ptr<Table>> RecordBatchFileReader::ToTable() {
  ARROW_ASSIGN_OR_RAISE(auto batches, ToRecordBatches());
  return Table::FromRecordBatches(schema(), std::move(batches));
}

Status Listener::OnEOS() { return Status::OK(); }

Status Listener::OnSchemaDecoded(std::shared_ptr<Schema> schema) { return Status::OK(); }

Status Listener::OnSchemaDecoded(std::shared_ptr<Schema> schema,
                                 std::shared_ptr<Schema> filtered_schema) {
  return OnSchemaDecoded(std::move(schema));
}

Status Listener::OnRecordBatchDecoded(std::shared_ptr<RecordBatch> record_batch) {
  return Status::NotImplemented("OnRecordBatchDecoded() callback isn't implemented");
}

Status Listener::OnRecordBatchWithMetadataDecoded(
    RecordBatchWithMetadata record_batch_with_metadata) {
  return OnRecordBatchDecoded(std::move(record_batch_with_metadata.batch));
}

class StreamDecoder::StreamDecoderImpl : public StreamDecoderInternal {
 public:
  explicit StreamDecoderImpl(std::shared_ptr<Listener> listener, IpcReadOptions options)
      : StreamDecoderInternal(std::move(listener), options),
        message_decoder_(std::shared_ptr<StreamDecoderImpl>(this, [](void*) {}),
                         options.memory_pool) {}

  Status Consume(const uint8_t* data, int64_t size) {
    return message_decoder_.Consume(data, size);
  }

  Status Consume(std::shared_ptr<Buffer> buffer) {
    return message_decoder_.Consume(std::move(buffer));
  }

  int64_t next_required_size() const { return message_decoder_.next_required_size(); }

  const MessageDecoder* message_decoder() const { return &message_decoder_; }

 private:
  MessageDecoder message_decoder_;
};

StreamDecoder::StreamDecoder(std::shared_ptr<Listener> listener, IpcReadOptions options) {
  impl_ = std::make_unique<StreamDecoderImpl>(std::move(listener), options);
}

StreamDecoder::~StreamDecoder() {}

Status StreamDecoder::Consume(const uint8_t* data, int64_t size) {
  while (size > 0) {
    const auto next_required_size = impl_->next_required_size();
    if (next_required_size == 0) {
      break;
    }
    if (size < next_required_size) {
      break;
    }
    ARROW_RETURN_NOT_OK(impl_->Consume(data, next_required_size));
    data += next_required_size;
    size -= next_required_size;
  }
  if (size > 0) {
    return impl_->Consume(data, size);
  } else {
    return arrow::Status::OK();
  }
}

Status StreamDecoder::Consume(std::shared_ptr<Buffer> buffer) {
  if (buffer->size() == 0) {
    return arrow::Status::OK();
  }
  if (impl_->next_required_size() == 0 || buffer->size() <= impl_->next_required_size()) {
    return impl_->Consume(std::move(buffer));
  } else {
    int64_t offset = 0;
    while (true) {
      const auto next_required_size = impl_->next_required_size();
      if (next_required_size == 0) {
        break;
      }
      if (buffer->size() - offset <= next_required_size) {
        break;
      }
      if (buffer->is_cpu()) {
        switch (impl_->message_decoder()->state()) {
          case MessageDecoder::State::INITIAL:
          case MessageDecoder::State::METADATA_LENGTH:
            // We don't need to pass a sliced buffer because
            // MessageDecoder doesn't keep reference of the given
            // buffer on these states.
            ARROW_RETURN_NOT_OK(
                impl_->Consume(buffer->data() + offset, next_required_size));
            break;
          default:
            ARROW_RETURN_NOT_OK(
                impl_->Consume(SliceBuffer(buffer, offset, next_required_size)));
            break;
        }
      } else {
        ARROW_RETURN_NOT_OK(
            impl_->Consume(SliceBuffer(buffer, offset, next_required_size)));
      }
      offset += next_required_size;
    }
    if (buffer->size() - offset == 0) {
      return arrow::Status::OK();
    } else if (offset == 0) {
      return impl_->Consume(std::move(buffer));
    } else {
      return impl_->Consume(SliceBuffer(std::move(buffer), offset));
    }
  }
}

Status StreamDecoder::Reset() {
  impl_ = std::make_unique<StreamDecoderImpl>(impl_->listener(), impl_->options());
  return Status::OK();
}

std::shared_ptr<Schema> StreamDecoder::schema() const { return impl_->schema(); }

int64_t StreamDecoder::next_required_size() const { return impl_->next_required_size(); }

ReadStats StreamDecoder::stats() const { return impl_->stats(); }

Result<std::shared_ptr<Schema>> ReadSchema(io::InputStream* stream,
                                           DictionaryMemo* dictionary_memo) {
  std::unique_ptr<MessageReader> reader = MessageReader::Open(stream);
  ARROW_ASSIGN_OR_RAISE(std::unique_ptr<Message> message, reader->ReadNextMessage());
  if (!message) {
    return Status::Invalid("Tried reading schema message, was null or length 0");
  }
  CHECK_MESSAGE_TYPE(MessageType::SCHEMA, message->type());
  return ReadSchema(*message, dictionary_memo);
}

Result<std::shared_ptr<Schema>> ReadSchema(const Message& message,
                                           DictionaryMemo* dictionary_memo) {
  std::shared_ptr<Schema> result;
  RETURN_NOT_OK(internal::GetSchema(message.header(), dictionary_memo, &result));
  return result;
}

Result<std::shared_ptr<Tensor>> ReadTensor(io::InputStream* file) {
  std::unique_ptr<Message> message;
  RETURN_NOT_OK(ReadContiguousPayload(file, &message));
  return ReadTensor(*message);
}

Result<std::shared_ptr<Tensor>> ReadTensor(const Message& message) {
  std::shared_ptr<DataType> type;
  std::vector<int64_t> shape;
  std::vector<int64_t> strides;
  std::vector<std::string> dim_names;
  CHECK_HAS_BODY(message);
  RETURN_NOT_OK(internal::GetTensorMetadata(*message.metadata(), &type, &shape, &strides,
                                            &dim_names));
  return Tensor::Make(type, message.body(), shape, strides, dim_names);
}

namespace {

Result<std::shared_ptr<SparseIndex>> ReadSparseCOOIndex(
    const flatbuf::SparseTensor* sparse_tensor, const std::vector<int64_t>& shape,
    int64_t non_zero_length, io::RandomAccessFile* file) {
  auto* sparse_index = sparse_tensor->sparseIndex_as_SparseTensorIndexCOO();
  const auto ndim = static_cast<int64_t>(shape.size());

  std::shared_ptr<DataType> indices_type;
  RETURN_NOT_OK(internal::GetSparseCOOIndexMetadata(sparse_index, &indices_type));
  const int64_t indices_elsize = indices_type->byte_width();

  auto* indices_buffer = sparse_index->indicesBuffer();
  ARROW_ASSIGN_OR_RAISE(auto indices_data,
                        file->ReadAt(indices_buffer->offset(), indices_buffer->length()));
  std::vector<int64_t> indices_shape({non_zero_length, ndim});
  auto* indices_strides = sparse_index->indicesStrides();
  std::vector<int64_t> strides(2);
  if (indices_strides && indices_strides->size() > 0) {
    if (indices_strides->size() != 2) {
      return Status::Invalid("Wrong size for indicesStrides in SparseCOOIndex");
    }
    strides[0] = indices_strides->Get(0);
    strides[1] = indices_strides->Get(1);
  } else {
    // Row-major by default
    strides[0] = indices_elsize * ndim;
    strides[1] = indices_elsize;
  }
  return SparseCOOIndex::Make(
      std::make_shared<Tensor>(indices_type, indices_data, indices_shape, strides),
      sparse_index->isCanonical());
}

Result<std::shared_ptr<SparseIndex>> ReadSparseCSXIndex(
    const flatbuf::SparseTensor* sparse_tensor, const std::vector<int64_t>& shape,
    int64_t non_zero_length, io::RandomAccessFile* file) {
  if (shape.size() != 2) {
    return Status::Invalid("Invalid shape length for a sparse matrix");
  }

  auto* sparse_index = sparse_tensor->sparseIndex_as_SparseMatrixIndexCSX();

  std::shared_ptr<DataType> indptr_type, indices_type;
  RETURN_NOT_OK(
      internal::GetSparseCSXIndexMetadata(sparse_index, &indptr_type, &indices_type));
  const int indptr_byte_width = indptr_type->byte_width();

  auto* indptr_buffer = sparse_index->indptrBuffer();
  ARROW_ASSIGN_OR_RAISE(auto indptr_data,
                        file->ReadAt(indptr_buffer->offset(), indptr_buffer->length()));

  auto* indices_buffer = sparse_index->indicesBuffer();
  ARROW_ASSIGN_OR_RAISE(auto indices_data,
                        file->ReadAt(indices_buffer->offset(), indices_buffer->length()));

  std::vector<int64_t> indices_shape({non_zero_length});
  const auto indices_minimum_bytes = indices_shape[0] * indices_type->byte_width();
  if (indices_minimum_bytes > indices_buffer->length()) {
    return Status::Invalid("shape is inconsistent to the size of indices buffer");
  }

  switch (sparse_index->compressedAxis()) {
    case flatbuf::SparseMatrixCompressedAxis::SparseMatrixCompressedAxis_Row: {
      std::vector<int64_t> indptr_shape({shape[0] + 1});
      const int64_t indptr_minimum_bytes = indptr_shape[0] * indptr_byte_width;
      if (indptr_minimum_bytes > indptr_buffer->length()) {
        return Status::Invalid("shape is inconsistent to the size of indptr buffer");
      }
      return std::make_shared<SparseCSRIndex>(
          std::make_shared<Tensor>(indptr_type, indptr_data, indptr_shape),
          std::make_shared<Tensor>(indices_type, indices_data, indices_shape));
    }
    case flatbuf::SparseMatrixCompressedAxis::SparseMatrixCompressedAxis_Column: {
      std::vector<int64_t> indptr_shape({shape[1] + 1});
      const int64_t indptr_minimum_bytes = indptr_shape[0] * indptr_byte_width;
      if (indptr_minimum_bytes > indptr_buffer->length()) {
        return Status::Invalid("shape is inconsistent to the size of indptr buffer");
      }
      return std::make_shared<SparseCSCIndex>(
          std::make_shared<Tensor>(indptr_type, indptr_data, indptr_shape),
          std::make_shared<Tensor>(indices_type, indices_data, indices_shape));
    }
    default:
      return Status::Invalid("Invalid value of SparseMatrixCompressedAxis");
  }
}

Result<std::shared_ptr<SparseIndex>> ReadSparseCSFIndex(
    const flatbuf::SparseTensor* sparse_tensor, const std::vector<int64_t>& shape,
    io::RandomAccessFile* file) {
  auto* sparse_index = sparse_tensor->sparseIndex_as_SparseTensorIndexCSF();
  const auto ndim = static_cast<int64_t>(shape.size());
  auto* indptr_buffers = sparse_index->indptrBuffers();
  auto* indices_buffers = sparse_index->indicesBuffers();
  std::vector<std::shared_ptr<Buffer>> indptr_data(ndim - 1);
  std::vector<std::shared_ptr<Buffer>> indices_data(ndim);

  std::shared_ptr<DataType> indptr_type, indices_type;
  std::vector<int64_t> axis_order, indices_size;

  RETURN_NOT_OK(internal::GetSparseCSFIndexMetadata(
      sparse_index, &axis_order, &indices_size, &indptr_type, &indices_type));
  for (int i = 0; i < static_cast<int>(indptr_buffers->size()); ++i) {
    ARROW_ASSIGN_OR_RAISE(indptr_data[i], file->ReadAt(indptr_buffers->Get(i)->offset(),
                                                       indptr_buffers->Get(i)->length()));
  }
  for (int i = 0; i < static_cast<int>(indices_buffers->size()); ++i) {
    ARROW_ASSIGN_OR_RAISE(indices_data[i],
                          file->ReadAt(indices_buffers->Get(i)->offset(),
                                       indices_buffers->Get(i)->length()));
  }

  return SparseCSFIndex::Make(indptr_type, indices_type, indices_size, axis_order,
                              indptr_data, indices_data);
}

Result<std::shared_ptr<SparseTensor>> MakeSparseTensorWithSparseCOOIndex(
    const std::shared_ptr<DataType>& type, const std::vector<int64_t>& shape,
    const std::vector<std::string>& dim_names,
    const std::shared_ptr<SparseCOOIndex>& sparse_index, int64_t non_zero_length,
    const std::shared_ptr<Buffer>& data) {
  return SparseCOOTensor::Make(sparse_index, type, data, shape, dim_names);
}

Result<std::shared_ptr<SparseTensor>> MakeSparseTensorWithSparseCSRIndex(
    const std::shared_ptr<DataType>& type, const std::vector<int64_t>& shape,
    const std::vector<std::string>& dim_names,
    const std::shared_ptr<SparseCSRIndex>& sparse_index, int64_t non_zero_length,
    const std::shared_ptr<Buffer>& data) {
  return SparseCSRMatrix::Make(sparse_index, type, data, shape, dim_names);
}

Result<std::shared_ptr<SparseTensor>> MakeSparseTensorWithSparseCSCIndex(
    const std::shared_ptr<DataType>& type, const std::vector<int64_t>& shape,
    const std::vector<std::string>& dim_names,
    const std::shared_ptr<SparseCSCIndex>& sparse_index, int64_t non_zero_length,
    const std::shared_ptr<Buffer>& data) {
  return SparseCSCMatrix::Make(sparse_index, type, data, shape, dim_names);
}

Result<std::shared_ptr<SparseTensor>> MakeSparseTensorWithSparseCSFIndex(
    const std::shared_ptr<DataType>& type, const std::vector<int64_t>& shape,
    const std::vector<std::string>& dim_names,
    const std::shared_ptr<SparseCSFIndex>& sparse_index,
    const std::shared_ptr<Buffer>& data) {
  return SparseCSFTensor::Make(sparse_index, type, data, shape, dim_names);
}

Status ReadSparseTensorMetadata(const Buffer& metadata,
                                std::shared_ptr<DataType>* out_type,
                                std::vector<int64_t>* out_shape,
                                std::vector<std::string>* out_dim_names,
                                int64_t* out_non_zero_length,
                                SparseTensorFormat::type* out_format_id,
                                const flatbuf::SparseTensor** out_fb_sparse_tensor,
                                const flatbuf::Buffer** out_buffer) {
  RETURN_NOT_OK(internal::GetSparseTensorMetadata(
      metadata, out_type, out_shape, out_dim_names, out_non_zero_length, out_format_id));

  const flatbuf::Message* message = nullptr;
  RETURN_NOT_OK(internal::VerifyMessage(metadata.data(), metadata.size(), &message));

  auto sparse_tensor = message->header_as_SparseTensor();
  if (sparse_tensor == nullptr) {
    return Status::IOError(
        "Header-type of flatbuffer-encoded Message is not SparseTensor.");
  }
  *out_fb_sparse_tensor = sparse_tensor;

  auto buffer = sparse_tensor->data();
  if (!bit_util::IsMultipleOf8(buffer->offset())) {
    return Status::Invalid(
        "Buffer of sparse index data did not start on 8-byte aligned offset: ",
        buffer->offset());
  }
  *out_buffer = buffer;

  return Status::OK();
}

}  // namespace

namespace internal {

namespace {

Result<size_t> GetSparseTensorBodyBufferCount(SparseTensorFormat::type format_id,
                                              const size_t ndim) {
  switch (format_id) {
    case SparseTensorFormat::COO:
      return 2;

    case SparseTensorFormat::CSR:
      return 3;

    case SparseTensorFormat::CSC:
      return 3;

    case SparseTensorFormat::CSF:
      return 2 * ndim;

    default:
      return Status::Invalid("Unrecognized sparse tensor format");
  }
}

Status CheckSparseTensorBodyBufferCount(const IpcPayload& payload,
                                        SparseTensorFormat::type sparse_tensor_format_id,
                                        const size_t ndim) {
  size_t expected_body_buffer_count = 0;
  ARROW_ASSIGN_OR_RAISE(expected_body_buffer_count,
                        GetSparseTensorBodyBufferCount(sparse_tensor_format_id, ndim));
  if (payload.body_buffers.size() != expected_body_buffer_count) {
    return Status::Invalid("Invalid body buffer count for a sparse tensor");
  }

  return Status::OK();
}

}  // namespace

Result<size_t> ReadSparseTensorBodyBufferCount(const Buffer& metadata) {
  SparseTensorFormat::type format_id{};
  std::vector<int64_t> shape;

  RETURN_NOT_OK(internal::GetSparseTensorMetadata(metadata, nullptr, &shape, nullptr,
                                                  nullptr, &format_id));

  return GetSparseTensorBodyBufferCount(format_id, static_cast<size_t>(shape.size()));
}

Result<std::shared_ptr<SparseTensor>> ReadSparseTensorPayload(const IpcPayload& payload) {
  std::shared_ptr<DataType> type;
  std::vector<int64_t> shape;
  std::vector<std::string> dim_names;
  int64_t non_zero_length;
  SparseTensorFormat::type sparse_tensor_format_id;
  const flatbuf::SparseTensor* sparse_tensor;
  const flatbuf::Buffer* buffer;

  RETURN_NOT_OK(ReadSparseTensorMetadata(*payload.metadata, &type, &shape, &dim_names,
                                         &non_zero_length, &sparse_tensor_format_id,
                                         &sparse_tensor, &buffer));

  RETURN_NOT_OK(CheckSparseTensorBodyBufferCount(payload, sparse_tensor_format_id,
                                                 static_cast<size_t>(shape.size())));

  switch (sparse_tensor_format_id) {
    case SparseTensorFormat::COO: {
      std::shared_ptr<SparseCOOIndex> sparse_index;
      std::shared_ptr<DataType> indices_type;
      RETURN_NOT_OK(internal::GetSparseCOOIndexMetadata(
          sparse_tensor->sparseIndex_as_SparseTensorIndexCOO(), &indices_type));
      ARROW_ASSIGN_OR_RAISE(sparse_index,
                            SparseCOOIndex::Make(indices_type, shape, non_zero_length,
                                                 payload.body_buffers[0]));
      return MakeSparseTensorWithSparseCOOIndex(type, shape, dim_names, sparse_index,
                                                non_zero_length, payload.body_buffers[1]);
    }
    case SparseTensorFormat::CSR: {
      std::shared_ptr<SparseCSRIndex> sparse_index;
      std::shared_ptr<DataType> indptr_type;
      std::shared_ptr<DataType> indices_type;
      RETURN_NOT_OK(internal::GetSparseCSXIndexMetadata(
          sparse_tensor->sparseIndex_as_SparseMatrixIndexCSX(), &indptr_type,
          &indices_type));
      ARROW_CHECK_EQ(indptr_type, indices_type);
      ARROW_ASSIGN_OR_RAISE(
          sparse_index,
          SparseCSRIndex::Make(indices_type, shape, non_zero_length,
                               payload.body_buffers[0], payload.body_buffers[1]));
      return MakeSparseTensorWithSparseCSRIndex(type, shape, dim_names, sparse_index,
                                                non_zero_length, payload.body_buffers[2]);
    }
    case SparseTensorFormat::CSC: {
      std::shared_ptr<SparseCSCIndex> sparse_index;
      std::shared_ptr<DataType> indptr_type;
      std::shared_ptr<DataType> indices_type;
      RETURN_NOT_OK(internal::GetSparseCSXIndexMetadata(
          sparse_tensor->sparseIndex_as_SparseMatrixIndexCSX(), &indptr_type,
          &indices_type));
      ARROW_CHECK_EQ(indptr_type, indices_type);
      ARROW_ASSIGN_OR_RAISE(
          sparse_index,
          SparseCSCIndex::Make(indices_type, shape, non_zero_length,
                               payload.body_buffers[0], payload.body_buffers[1]));
      return MakeSparseTensorWithSparseCSCIndex(type, shape, dim_names, sparse_index,
                                                non_zero_length, payload.body_buffers[2]);
    }
    case SparseTensorFormat::CSF: {
      std::shared_ptr<SparseCSFIndex> sparse_index;
      std::shared_ptr<DataType> indptr_type, indices_type;
      std::vector<int64_t> axis_order, indices_size;

      RETURN_NOT_OK(internal::GetSparseCSFIndexMetadata(
          sparse_tensor->sparseIndex_as_SparseTensorIndexCSF(), &axis_order,
          &indices_size, &indptr_type, &indices_type));
      ARROW_CHECK_EQ(indptr_type, indices_type);

      const int64_t ndim = shape.size();
      std::vector<std::shared_ptr<Buffer>> indptr_data(ndim - 1);
      std::vector<std::shared_ptr<Buffer>> indices_data(ndim);

      for (int64_t i = 0; i < ndim - 1; ++i) {
        indptr_data[i] = payload.body_buffers[i];
      }
      for (int64_t i = 0; i < ndim; ++i) {
        indices_data[i] = payload.body_buffers[i + ndim - 1];
      }

      ARROW_ASSIGN_OR_RAISE(sparse_index,
                            SparseCSFIndex::Make(indptr_type, indices_type, indices_size,
                                                 axis_order, indptr_data, indices_data));
      return MakeSparseTensorWithSparseCSFIndex(type, shape, dim_names, sparse_index,
                                                payload.body_buffers[2 * ndim - 1]);
    }
    default:
      return Status::Invalid("Unsupported sparse index format");
  }
}

}  // namespace internal

namespace {

Result<std::shared_ptr<SparseTensor>> ReadSparseTensor(const Buffer& metadata,
                                                       io::RandomAccessFile* file) {
  std::shared_ptr<DataType> type;
  std::vector<int64_t> shape;
  std::vector<std::string> dim_names;
  int64_t non_zero_length;
  SparseTensorFormat::type sparse_tensor_format_id;
  const flatbuf::SparseTensor* sparse_tensor;
  const flatbuf::Buffer* buffer;

  RETURN_NOT_OK(ReadSparseTensorMetadata(metadata, &type, &shape, &dim_names,
                                         &non_zero_length, &sparse_tensor_format_id,
                                         &sparse_tensor, &buffer));

  ARROW_ASSIGN_OR_RAISE(auto data, file->ReadAt(buffer->offset(), buffer->length()));

  std::shared_ptr<SparseIndex> sparse_index;
  switch (sparse_tensor_format_id) {
    case SparseTensorFormat::COO: {
      ARROW_ASSIGN_OR_RAISE(
          sparse_index, ReadSparseCOOIndex(sparse_tensor, shape, non_zero_length, file));
      return MakeSparseTensorWithSparseCOOIndex(
          type, shape, dim_names, checked_pointer_cast<SparseCOOIndex>(sparse_index),
          non_zero_length, data);
    }
    case SparseTensorFormat::CSR: {
      ARROW_ASSIGN_OR_RAISE(
          sparse_index, ReadSparseCSXIndex(sparse_tensor, shape, non_zero_length, file));
      return MakeSparseTensorWithSparseCSRIndex(
          type, shape, dim_names, checked_pointer_cast<SparseCSRIndex>(sparse_index),
          non_zero_length, data);
    }
    case SparseTensorFormat::CSC: {
      ARROW_ASSIGN_OR_RAISE(
          sparse_index, ReadSparseCSXIndex(sparse_tensor, shape, non_zero_length, file));
      return MakeSparseTensorWithSparseCSCIndex(
          type, shape, dim_names, checked_pointer_cast<SparseCSCIndex>(sparse_index),
          non_zero_length, data);
    }
    case SparseTensorFormat::CSF: {
      ARROW_ASSIGN_OR_RAISE(sparse_index, ReadSparseCSFIndex(sparse_tensor, shape, file));
      return MakeSparseTensorWithSparseCSFIndex(
          type, shape, dim_names, checked_pointer_cast<SparseCSFIndex>(sparse_index),
          data);
    }
    default:
      return Status::Invalid("Unsupported sparse index format");
  }
}

}  // namespace

Result<std::shared_ptr<SparseTensor>> ReadSparseTensor(const Message& message) {
  CHECK_HAS_BODY(message);
  ARROW_ASSIGN_OR_RAISE(auto reader, Buffer::GetReader(message.body()));
  return ReadSparseTensor(*message.metadata(), reader.get());
}

Result<std::shared_ptr<SparseTensor>> ReadSparseTensor(io::InputStream* file) {
  std::unique_ptr<Message> message;
  RETURN_NOT_OK(ReadContiguousPayload(file, &message));
  CHECK_MESSAGE_TYPE(MessageType::SPARSE_TENSOR, message->type());
  CHECK_HAS_BODY(*message);
  ARROW_ASSIGN_OR_RAISE(auto reader, Buffer::GetReader(message->body()));
  return ReadSparseTensor(*message->metadata(), reader.get());
}

///////////////////////////////////////////////////////////////////////////
// Helpers for fuzzing

namespace internal {
namespace {

Status ValidateFuzzBatch(const RecordBatch& batch) {
  auto st = batch.ValidateFull();
  if (st.ok()) {
    // If the batch is valid, printing should succeed
    batch.ToString();
  }
  return st;
}

Status ValidateFuzzBatch(const RecordBatchWithMetadata& batch) {
  if (batch.batch) {
    RETURN_NOT_OK(ValidateFuzzBatch(*batch.batch));
  }
  // XXX do something with custom metadata?
  return Status::OK();
}

Status CompareFuzzBatches(const RecordBatchWithMetadata& left,
                          const RecordBatchWithMetadata& right) {
  if ((left.custom_metadata != nullptr) != (right.custom_metadata != nullptr) ||
      (left.custom_metadata && !left.custom_metadata->Equals(*right.custom_metadata))) {
    std::stringstream ss;
    auto print_metadata = [&](const RecordBatchWithMetadata& batch) {
      if (batch.custom_metadata) {
        ss << batch.custom_metadata->ToString();
      } else {
        ss << "nullptr";
      }
    };
    ss << "Custom metadata unequal: left = ";
    print_metadata(left);
    ss << "\nright = ";
    print_metadata(right);
    return Status::Invalid(ss.str());
  }

  bool ok = true;
  if ((left.batch != nullptr) != (right.batch != nullptr)) {
    ok = false;
  } else if (left.batch) {
    ok &= left.batch->Equals(*right.batch, EqualOptions{}.nans_equal(true));
  }
  if (ok) {
    return Status::OK();
  }
  std::stringstream ss;
  PrettyPrintOptions options{};
  options.show_field_metadata = true;
  options.show_schema_metadata = true;
  auto print_batch = [&](const RecordBatchWithMetadata& batch) {
    if (batch.batch) {
      ss << "\n";
      ARROW_UNUSED(PrettyPrint(*batch.batch, options, &ss));
    } else {
      ss << "nullptr";
    }
  };
  ss << "Batches unequal: left = ";
  print_batch(left);
  ss << "\nright = ";
  print_batch(right);
  return Status::Invalid(ss.str());
}

Status CompareFuzzBatches(const std::vector<RecordBatchWithMetadata>& left,
                          const std::vector<RecordBatchWithMetadata>& right) {
  if (left.size() != right.size()) {
    return Status::Invalid("Not the same number of batches");
  }
  for (size_t i = 0; i < left.size(); ++i) {
    RETURN_NOT_OK(CompareFuzzBatches(left[i], right[i]));
  }
  return Status::OK();
}

IpcReadOptions FuzzingOptions() {
  IpcReadOptions options;
  options.memory_pool = ::arrow::internal::fuzzing_memory_pool();
  return options;
}

}  // namespace

Status FuzzIpcStream(const uint8_t* data, int64_t size) {
  auto buffer = std::make_shared<Buffer>(data, size);
  io::BufferReader buffer_reader(buffer);

  std::shared_ptr<RecordBatchReader> batch_reader;
  ARROW_ASSIGN_OR_RAISE(batch_reader,
                        RecordBatchStreamReader::Open(&buffer_reader, FuzzingOptions()));
  Status st;

  while (true) {
    ARROW_ASSIGN_OR_RAISE(auto batch, batch_reader->ReadNext());
    if (!batch.batch && !batch.custom_metadata) {
      // EOS
      break;
    }
    st &= ValidateFuzzBatch(batch);
  }

  return st;
}

Status FuzzIpcFile(const uint8_t* data, int64_t size) {
  auto buffer = std::make_shared<Buffer>(data, size);

  Status final_status;

  struct IpcReadResult {
    std::shared_ptr<Schema> schema;
    std::vector<RecordBatchWithMetadata> batches = {};
  };

  // Try to read the IPC file as a stream to compare the results (differential fuzzing)
  auto do_stream_read = [&]() -> Result<IpcReadResult> {
    io::BufferReader buffer_reader(buffer);
    // Skip magic bytes at the beginning
    RETURN_NOT_OK(
        buffer_reader.Advance(bit_util::RoundUpToMultipleOf8(kArrowMagicBytes.length())));
    ARROW_ASSIGN_OR_RAISE(auto batch_reader, RecordBatchStreamReader::Open(
                                                 &buffer_reader, FuzzingOptions()));

    std::vector<RecordBatchWithMetadata> batches;
    while (true) {
      ARROW_ASSIGN_OR_RAISE(auto batch, batch_reader->ReadNext());
      if (!batch.batch && !batch.custom_metadata) {
        // EOS
        break;
      }
      RETURN_NOT_OK(ValidateFuzzBatch(batch));
      batches.push_back(batch);
    }
    return IpcReadResult{batch_reader->schema(), batches};
  };

  auto do_file_read = [&](bool pre_buffer) -> Result<IpcReadResult> {
    io::BufferReader buffer_reader(buffer);
    ARROW_ASSIGN_OR_RAISE(auto batch_reader,
                          RecordBatchFileReader::Open(&buffer_reader, FuzzingOptions()));
    if (pre_buffer) {
      // Pre-buffer all record batches
      RETURN_NOT_OK(batch_reader->PreBufferMetadata(/*indices=*/{}));
    }

    IpcReadResult result = {.schema = batch_reader->schema()};
    const int n_batches = batch_reader->num_record_batches();
    // Delay error return until the end, as we want to access all record batches
    Status st;
    for (int i = 0; i < n_batches; ++i) {
      RecordBatchWithMetadata batch;
      st &= batch_reader->ReadRecordBatchWithCustomMetadata(i).Value(&batch);
      st &= ValidateFuzzBatch(batch);
      result.batches.push_back(batch);
    }
    RETURN_NOT_OK(st);
    return result;
  };

  // Persistent read result for differential fuzzing.
  std::optional<IpcReadResult> maybe_read_result;

  auto compare_result = [&](IpcReadResult new_result) {
    // Differential fuzzing: compare results read using different APIs or options
    // with one another.
    if (maybe_read_result.has_value()) {
      ARROW_CHECK_OK(CompareFuzzBatches(maybe_read_result->batches, new_result.batches));
    } else {
      maybe_read_result = std::move(new_result);
    }
  };

  for (const bool pre_buffer : {false, true}) {
    auto maybe_file_result = do_file_read(pre_buffer);
    final_status &= maybe_file_result.status();
    if (maybe_file_result.ok()) {
      compare_result(*std::move(maybe_file_result));
    }
  }

  if (maybe_read_result.has_value()) {
    // IPC file read successful: compare results with IPC stream reader,
    // if possible.
    // NOTE: some valid IPC files may not be readable as IPC streams,
    // for example because of excess spacing between IPC messages.
    // A regular IPC file writer would not produce them, but fuzzing might.
    auto maybe_stream_result = do_stream_read();
    final_status &= maybe_stream_result.status();
    if (maybe_stream_result.ok()) {
      if (maybe_read_result->schema->Equals(maybe_stream_result->schema,
                                            /*check_metadata=*/true)) {
        // XXX: in some rare cases, an IPC file might read unequal to the enclosed
        // IPC stream, for example if the footer skips some batches or orders the
        // batches differently. We should revisit this if the fuzzer generates such
        // files.
        compare_result(*maybe_stream_result);
      } else {
        // The fuzzer might have mutated the schema definition that is duplicated
        // in the IPC file footer, in which case the comparison above would fail.
        final_status &= Status::TypeError(
            "Schema mismatch between IPC stream and IPC file footer, skipping "
            "comparison");
      }
    }
  }
  return final_status;
}

Status FuzzIpcTensorStream(const uint8_t* data, int64_t size) {
  auto buffer = std::make_shared<Buffer>(data, size);
  io::BufferReader buffer_reader(buffer);

  std::shared_ptr<Tensor> tensor;

  while (true) {
    ARROW_ASSIGN_OR_RAISE(tensor, ReadTensor(&buffer_reader));
    if (tensor == nullptr) {
      break;
    }
    RETURN_NOT_OK(tensor->Validate());
  }

  return Status::OK();
}

Result<int64_t> IoRecordedRandomAccessFile::GetSize() { return file_size_; }

Result<int64_t> IoRecordedRandomAccessFile::ReadAt(int64_t position, int64_t nbytes,
                                                   void* out) {
  auto num_bytes_read = std::min(file_size_, position + nbytes) - position;

  if (!read_ranges_.empty() &&
      position == read_ranges_.back().offset + read_ranges_.back().length) {
    // merge continuous IOs into one if possible
    read_ranges_.back().length += num_bytes_read;
  } else {
    // no real IO is performed, it is only saved into a vector for replaying later
    read_ranges_.emplace_back(io::ReadRange{position, num_bytes_read});
  }
  return num_bytes_read;
}

Result<std::shared_ptr<Buffer>> IoRecordedRandomAccessFile::ReadAt(int64_t position,
                                                                   int64_t nbytes) {
  std::shared_ptr<Buffer> out;
  auto result = ReadAt(position, nbytes, &out);
  return out;
}

Status IoRecordedRandomAccessFile::Close() {
  closed_ = true;
  return Status::OK();
}

Status IoRecordedRandomAccessFile::Abort() { return Status::OK(); }

Result<int64_t> IoRecordedRandomAccessFile::Tell() const { return position_; }

bool IoRecordedRandomAccessFile::closed() const { return closed_; }

Status IoRecordedRandomAccessFile::Seek(int64_t position) {
  position_ = position;
  return Status::OK();
}

Result<int64_t> IoRecordedRandomAccessFile::Read(int64_t nbytes, void* out) {
  ARROW_ASSIGN_OR_RAISE(int64_t bytes_read, ReadAt(position_, nbytes, out));
  position_ += bytes_read;
  return bytes_read;
}

Result<std::shared_ptr<Buffer>> IoRecordedRandomAccessFile::Read(int64_t nbytes) {
  ARROW_ASSIGN_OR_RAISE(std::shared_ptr<Buffer> buffer, ReadAt(position_, nbytes));
  auto num_bytes_read = std::min(file_size_, position_ + nbytes) - position_;
  position_ += num_bytes_read;
  return buffer;
}

const io::IOContext& IoRecordedRandomAccessFile::io_context() const {
  return io_context_;
}

const std::vector<io::ReadRange>& IoRecordedRandomAccessFile::GetReadRanges() const {
  return read_ranges_;
}

}  // namespace internal
}  // namespace ipc
}  // namespace arrow
