Repository: apache/arrow
File: cpp/src/arrow/flight/server.cc (C++)

Code before:

    // Return the expected schema payload.
    if (payload_deque_.empty()) {
      return Status::UnknownError("No schema payload generated");
    }
    *payload = std::move(payload_deque_.front());
    payload_deque_.pop_front();
    return Status::OK();
  }

  Status Next(FlightPayload* payload) {
    // If we have previous payloads (dictionary messages or previous record batches)
    // we will return them before reading the next record batch.
    if (payload_deque_.empty()) {
      std::shared_ptr<RecordBatch> batch;
      RETURN_NOT_OK(reader_->ReadNext(&batch));
      if (!batch) {
        // End of stream
        if (writer_) {
          RETURN_NOT_OK(writer_->Close());
        }
        payload->ipc_message.metadata = nullptr;
        return Status::OK();
      }
      if (!writer_) {
        RETURN_NOT_OK(InitializeWriter());
      }
      // One WriteRecordBatch call might generate multiple payloads, so we
      // need to collect them in a deque.
      RETURN_NOT_OK(writer_->WriteRecordBatch(*batch));
    }

    // There must be at least one payload generated after WriteRecordBatch or
    // from previous calls to WriteRecordBatch.
    if (payload_deque_.empty()) {
      return Status::UnknownError("IPC writer didn't produce any payloads");
    }

    *payload = std::move(payload_deque_.front());
    payload_deque_.pop_front();
    return Status::OK();
  }

  Status Close() {
    if (writer_) {
      RETURN_NOT_OK(writer_->Close());
    }
    return reader_->Close();
  }

 private:

Review comment (bug):
I think in this specific scenario we might have to drop the first message after calling `writer_->WriteRecordBatch`, which will be a schema message, after creation of the `ipc::internal::OpenRecordBatchWriter` and we want to keep the rest of the messages which should be the expected `RecordBatch` for those cases.
I would have to debug but from what I understand that might be the cause of the current problem.


Code after:

    // Return the expected schema payload.
    if (payload_deque_.empty()) {
      return Status::UnknownError("No schema payload generated");
    }
    *payload = std::move(payload_deque_.front());
    payload_deque_.pop_front();
    return Status::OK();
  }

  Status Next(FlightPayload* payload) {
    // If we have previous payloads (dictionary messages or previous record batches)
    // we will return them before reading the next record batch.
    if (payload_deque_.empty()) {
      std::shared_ptr<RecordBatch> batch;
      RETURN_NOT_OK(reader_->ReadNext(&batch));
      if (!batch) {
        // End of stream
        if (writer_) {
          RETURN_NOT_OK(writer_->Close());
        }
        payload->ipc_message.metadata = nullptr;
        return Status::OK();
      }
      if (!writer_) {
        RETURN_NOT_OK(InitializeWriter());
        // If the writer has not been initialized yet, the first batch in the payload
        // queue is going to be a SCHEMA one. In this context, that is
        // unexpected, so drop it from the queue so that there is a RECORD_BATCH
        // message on the top (same as would be if the writer had been initialized
        // in GetSchemaPayload).
        if (payload_deque_.front().ipc_message.type == ipc::MessageType::SCHEMA) {
          payload_deque_.pop_front();
        }
      }
      // One WriteRecordBatch call might generate multiple payloads, so we
      // need to collect them in a deque.
      RETURN_NOT_OK(writer_->WriteRecordBatch(*batch));
    }

    // There must be at least one payload generated after WriteRecordBatch or
    // from previous calls to WriteRecordBatch.
    if (payload_deque_.empty()) {
      return Status::UnknownError("IPC writer didn't produce any payloads");
    }

    *payload = std::move(payload_deque_.front());
    payload_deque_.pop_front();
    return Status::OK();
  }