public:

TransportIpcMessageReader(

std::shared_ptr<internal::PeekableFlightDataReader> peekable_reader,

std::shared_ptr<MemoryManager> memory_manager,

std::shared_ptr<Buffer>* app_metadata)

: peekable_reader_(peekable_reader),

memory_manager_(std::move(memory_manager)),

app_metadata_(app_metadata) {}

::arrow::Result<std::unique_ptr<ipc::Message>> ReadNextMessage() override {

if (stream_finished_) return nullptr;

internal::FlightData* data;

peekable_reader_->Next(&data);

if (!data) {

stream_finished_ = true;

return nullptr;

}

if (data->body) {

ARROW_ASSIGN_OR_RAISE(data->body, Buffer::ViewOrCopy(data->body, memory_manager_));

}

*app_metadata_ = std::move(data->app_metadata);

return data->OpenMessage();

}

protected:

std::shared_ptr<internal::PeekableFlightDataReader> peekable_reader_;

std::shared_ptr<MemoryManager> memory_manager_;

// A reference to TransportDataStream.app_metadata_. That class

// can't access the app metadata because when it Peek()s the stream,

// it may be looking at a dictionary batch, not the record

// batch. Updating it here ensures the reader is always updated with

// the last metadata message read.

std::shared_ptr<Buffer>* app_metadata_;

bool stream_finished_ = false;

public:

TransportMessageReader(ServerDataStream* stream,

std::shared_ptr<MemoryManager> memory_manager)

: peekable_reader_(new internal::PeekableFlightDataReader(stream)),

memory_manager_(std::move(memory_manager)) {}

Status Init() {

// Peek the first message to get the descriptor.

internal::FlightData* data;

peekable_reader_->Peek(&data);

if (!data) {

return Status::IOError("Stream finished before first message sent");

}

if (!data->descriptor) {

return Status::IOError("Descriptor missing on first message");

}

descriptor_ = *data->descriptor;

// If there's a schema (=DoPut), also Open().

if (data->metadata) {

return EnsureDataStarted();

}

peekable_reader_->Next(&data);

return Status::OK();

}

const FlightDescriptor& descriptor() const override { return descriptor_; }

arrow::Result<std::shared_ptr<Schema>> GetSchema() override {

RETURN_NOT_OK(EnsureDataStarted());

return batch_reader_->schema();

}

arrow::Result<FlightStreamChunk> Next() override {

FlightStreamChunk out;

internal::FlightData* data = nullptr;

peekable_reader_->Peek(&data);

if (!data) {

out.app_metadata = nullptr;

out.data = nullptr;

return out;

}

if (!data->metadata) {

// Metadata-only (data->metadata is the IPC header)

out.app_metadata = data->app_metadata;

out.data = nullptr;

peekable_reader_->Next(&data);

return out;

}

if (!batch_reader_) {

RETURN_NOT_OK(EnsureDataStarted());

// re-peek here since EnsureDataStarted() advances the stream

return Next();

}

RETURN_NOT_OK(batch_reader_->ReadNext(&out.data));

out.app_metadata = std::move(app_metadata_);

return out;

}

private:

/// Ensure we are set up to read data.

Status EnsureDataStarted() {

if (!batch_reader_) {

// peek() until we find the first data message; discard metadata

if (!peekable_reader_->SkipToData()) {

return Status::IOError("Client never sent a data message");

}

auto message_reader =

std::unique_ptr<ipc::MessageReader>(new TransportIpcMessageReader(

peekable_reader_, memory_manager_, &app_metadata_));

ARROW_ASSIGN_OR_RAISE(

batch_reader_, ipc::RecordBatchStreamReader::Open(std::move(message_reader)));

}

return Status::OK();

}

FlightDescriptor descriptor_;

std::shared_ptr<internal::PeekableFlightDataReader> peekable_reader_;

std::shared_ptr<MemoryManager> memory_manager_;

std::shared_ptr<RecordBatchReader> batch_reader_;

std::shared_ptr<Buffer> app_metadata_;

public:

explicit TransportMessageWriter(ServerDataStream* stream)

: stream_(stream), ipc_options_(::arrow::ipc::IpcWriteOptions::Defaults()) {}

Status Begin(const std::shared_ptr<Schema>& schema,

const ipc::IpcWriteOptions& options) override {

if (started_) {

return Status::Invalid("This writer has already been started.");

}

started_ = true;

ipc_options_ = options;

RETURN_NOT_OK(mapper_.AddSchemaFields(*schema));

FlightPayload schema_payload;

RETURN_NOT_OK(ipc::GetSchemaPayload(*schema, ipc_options_, mapper_,

&schema_payload.ipc_message));

return WritePayload(schema_payload);

}

Status WriteRecordBatch(const RecordBatch& batch) override {

return WriteWithMetadata(batch, nullptr);

}

Status WriteMetadata(std::shared_ptr<Buffer> app_metadata) override {

FlightPayload payload{};

payload.app_metadata = app_metadata;

return WritePayload(payload);

}

Status WriteWithMetadata(const RecordBatch& batch,

std::shared_ptr<Buffer> app_metadata) override {

RETURN_NOT_OK(CheckStarted());

RETURN_NOT_OK(EnsureDictionariesWritten(batch));

FlightPayload payload{};

if (app_metadata) {

payload.app_metadata = app_metadata;

}

RETURN_NOT_OK(ipc::GetRecordBatchPayload(batch, ipc_options_, &payload.ipc_message));

RETURN_NOT_OK(WritePayload(payload));

++stats_.num_record_batches;

return Status::OK();

}

Status Close() override {

// It's fine to Close() without writing data

return Status::OK();

}

ipc::WriteStats stats() const override { return stats_; }

private:

Status WritePayload(const FlightPayload& payload) {

ARROW_ASSIGN_OR_RAISE(auto success, stream_->WriteData(payload));

if (!success) {

return MakeFlightError(FlightStatusCode::Internal,

"Could not write metadata to stream (client disconnect?)");

}

++stats_.num_messages;

return Status::OK();

}

Status CheckStarted() {

if (!started_) {

return Status::Invalid("This writer is not started. Call Begin() with a schema");

}

return Status::OK();

}

Status EnsureDictionariesWritten(const RecordBatch& batch) {

if (dictionaries_written_) {

return Status::OK();

}

dictionaries_written_ = true;

ARROW_ASSIGN_OR_RAISE(const auto dictionaries,

ipc::CollectDictionaries(batch, mapper_));

for (const auto& pair : dictionaries) {

FlightPayload payload{};

RETURN_NOT_OK(ipc::GetDictionaryPayload(pair.first, pair.second, ipc_options_,

&payload.ipc_message));

RETURN_NOT_OK(WritePayload(payload));

++stats_.num_dictionary_batches;

}

return Status::OK();

}

ServerDataStream* stream_;

::arrow::ipc::IpcWriteOptions ipc_options_;

ipc::DictionaryFieldMapper mapper_;

ipc::WriteStats stats_;

bool started_ = false;

bool dictionaries_written_ = false;

public:

explicit TransportMetadataWriter(ServerDataStream* stream) : stream_(stream) {}

Status WriteMetadata(const Buffer& buffer) override {

return stream_->WritePutMetadata(buffer);

}

private:

ServerDataStream* stream_;

ServerDataStream* stream) {

std::unique_ptr<FlightDataStream> data_stream;

RETURN_NOT_OK(base_->DoGet(context, ticket, &data_stream));

if (!data_stream) return Status::KeyError("No data in this flight");

// Write the schema as the first message in the stream

ARROW_ASSIGN_OR_RAISE(auto schema_payload, data_stream->GetSchemaPayload());

ARROW_ASSIGN_OR_RAISE(auto success, stream->WriteData(schema_payload));

// Connection terminated

if (!success) return Status::OK();

// Consume data stream and write out payloads

while (true) {

ARROW_ASSIGN_OR_RAISE(FlightPayload payload, data_stream->Next());

// End of stream

if (payload.ipc_message.metadata == nullptr) break;

ARROW_ASSIGN_OR_RAISE(auto success, stream->WriteData(payload));

// Connection terminated

if (!success) return Status::OK();

}

RETURN_NOT_OK(stream->WritesDone());

return data_stream->Close();

ServerDataStream* stream) {

std::unique_ptr<TransportMessageReader> reader(

new TransportMessageReader(stream, memory_manager_));

std::unique_ptr<FlightMetadataWriter> writer(new TransportMetadataWriter(stream));

RETURN_NOT_OK(reader->Init());

RETURN_NOT_OK(base_->DoPut(context, std::move(reader), std::move(writer)));

RETURN_NOT_OK(stream->WritesDone());

return Status::OK();

ServerDataStream* stream) {

std::unique_ptr<TransportMessageReader> reader(

new TransportMessageReader(stream, memory_manager_));

std::unique_ptr<FlightMessageWriter> writer(new TransportMessageWriter(stream));

RETURN_NOT_OK(reader->Init());

RETURN_NOT_OK(base_->DoExchange(context, std::move(reader), std::move(writer)));

RETURN_NOT_OK(stream->WritesDone());

return Status::OK();