You are viewing a plain text version of this content. The canonical link for it is here.
Posted to github@arrow.apache.org by "raulcd (via GitHub)" <gi...@apache.org> on 2025/07/15 17:01:16 UTC
[PR] GH-45055: [C++][Flight] Support dictionary replacement in DoGet [arrow]
raulcd opened a new pull request, #47115:
URL: https://github.com/apache/arrow/pull/47115
TBD, WIP
### Rationale for this change
### What changes are included in this PR?
### Are these changes tested?
### Are there any user-facing changes?
**This PR includes breaking changes to public APIs.** (If there are any breaking changes to public APIs, please explain which changes are breaking. If not, you can remove this.)
**This PR contains a "Critical Fix".** (If the changes fix either (a) a security vulnerability, (b) a bug that caused incorrect or invalid data to be produced, or (c) a bug that causes a crash (even when the API contract is upheld), please provide explanation. If not, you can remove this.)
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] GH-45055: [C++][Flight] Update Flight Server RecordBatchStreamImpl to reuse ipc::RecordBatchWriter with custom IpcPayloadWriter instead of manually generating FlightPayload [arrow]
Posted by "lidavidm (via GitHub)" <gi...@apache.org>.
lidavidm commented on code in PR #47115:
URL: https://github.com/apache/arrow/pull/47115#discussion_r2227169956
##########
cpp/src/arrow/flight/sql/example/sqlite_tables_schema_batch_reader.cc:
##########
Review Comment:
Thanks for tracking this down!
##########
cpp/src/arrow/flight/server.cc:
##########
@@ -275,78 +277,99 @@ Status FlightServerBase::GetSchema(const ServerCallContext& context,
class RecordBatchStream::RecordBatchStreamImpl {
public:
- // Stages of the stream when producing payloads
- enum class Stage {
- NEW, // The stream has been created, but Next has not been called yet
- DICTIONARY, // Dictionaries have been collected, and are being sent
- RECORD_BATCH // Initial have been sent
- };
-
RecordBatchStreamImpl(const std::shared_ptr<RecordBatchReader>& reader,
const ipc::IpcWriteOptions& options)
- : reader_(reader), mapper_(*reader_->schema()), ipc_options_(options) {}
+ : reader_(reader), options_(options) {}
std::shared_ptr<Schema> schema() { return reader_->schema(); }
Status GetSchemaPayload(FlightPayload* payload) {
- return ipc::GetSchemaPayload(*reader_->schema(), ipc_options_, mapper_,
- &payload->ipc_message);
+ if (!writer_) {
+ // Create the IPC writer on first call
+ auto payload_writer =
+ std::make_unique<ServerRecordBatchPayloadWriter>(&payload_list_);
+ ARROW_ASSIGN_OR_RAISE(
+ writer_, ipc::internal::OpenRecordBatchWriter(std::move(payload_writer),
+ reader_->schema(), options_));
+ }
+
+ // Return the expected schema payload.
+ if (payload_list_.empty()) {
+ return Status::UnknownError("No schema payload generated");
+ }
+ *payload = std::move(payload_list_.front());
+ payload_list_.pop_front();
+ return Status::OK();
}
Status Next(FlightPayload* payload) {
- if (stage_ == Stage::NEW) {
- RETURN_NOT_OK(reader_->ReadNext(¤t_batch_));
- if (!current_batch_) {
- // Signal that iteration is over
+ // If we have previous payloads (dictionary messages or previous record batches)
+ // we will return them before reading the next record batch.
+ if (payload_list_.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();
}
- ARROW_ASSIGN_OR_RAISE(dictionaries_,
- ipc::CollectDictionaries(*current_batch_, mapper_));
- stage_ = Stage::DICTIONARY;
+ if (!writer_) {
+ return Status::UnknownError(
+ "Writer should be initialized before reading Next batches");
+ }
+ // One WriteRecordBatch call might generate multiple payloads, so we
+ // need to collect them in a list.
+ RETURN_NOT_OK(writer_->WriteRecordBatch(*batch));
}
- if (stage_ == Stage::DICTIONARY) {
- if (dictionary_index_ == static_cast<int>(dictionaries_.size())) {
- stage_ = Stage::RECORD_BATCH;
- return ipc::GetRecordBatchPayload(*current_batch_, ipc_options_,
- &payload->ipc_message);
- } else {
- return GetNextDictionary(payload);
- }
+ // There must be at least one payload generated after WriteRecordBatch or
+ // from previous calls to WriteRecordBatch.
+ if (payload_list_.empty()) {
+ return Status::UnknownError("IPC writer didn't produce any payloads");
}
- RETURN_NOT_OK(reader_->ReadNext(¤t_batch_));
+ *payload = std::move(payload_list_.front());
+ payload_list_.pop_front();
+ return Status::OK();
+ }
- // TODO(ARROW-10787): Delta dictionaries
- if (!current_batch_) {
- // Signal that iteration is over
- payload->ipc_message.metadata = nullptr;
- return Status::OK();
- } else {
- return ipc::GetRecordBatchPayload(*current_batch_, ipc_options_,
- &payload->ipc_message);
+ Status Close() {
+ if (writer_) {
+ RETURN_NOT_OK(writer_->Close());
}
+ return reader_->Close();
}
- Status Close() { return reader_->Close(); }
-
private:
- Status GetNextDictionary(FlightPayload* payload) {
- const auto& it = dictionaries_[dictionary_index_++];
- return ipc::GetDictionaryPayload(it.first, it.second, ipc_options_,
- &payload->ipc_message);
- }
+ // Simple payload writer that uses a list of generated payloads.
+ class ServerRecordBatchPayloadWriter : public ipc::internal::IpcPayloadWriter {
+ public:
+ explicit ServerRecordBatchPayloadWriter(std::list<FlightPayload>* payload_list)
+ : payload_list_(payload_list) {}
- Stage stage_ = Stage::NEW;
- std::shared_ptr<RecordBatchReader> reader_;
- ipc::DictionaryFieldMapper mapper_;
- ipc::IpcWriteOptions ipc_options_;
- std::shared_ptr<RecordBatch> current_batch_;
- std::vector<std::pair<int64_t, std::shared_ptr<Array>>> dictionaries_;
+ Status Start() override { return Status::OK(); }
- // Index of next dictionary to send
- int dictionary_index_ = 0;
+ Status WritePayload(const ipc::IpcPayload& ipc_payload) override {
+ FlightPayload payload;
+ payload.ipc_message = ipc_payload;
+
+ payload_list_->push_back(std::move(payload));
+ return Status::OK();
+ }
+
+ Status Close() override { return Status::OK(); }
+
+ private:
+ std::list<FlightPayload>* payload_list_;
+ };
+
+ std::shared_ptr<RecordBatchReader> reader_;
+ ipc::IpcWriteOptions options_;
+ std::unique_ptr<ipc::RecordBatchWriter> writer_;
+ std::list<FlightPayload> payload_list_;
Review Comment:
Ah sorry one last thing: I would prefer `std::deque` (deque) for a double-ended queue over `std::list` (linked list)
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] GH-45055: [C++][Flight] Update Flight Server RecordBatchStreamImpl to reuse ipc::RecordBatchWriter with custom IpcPayloadWriter instead of manually generating FlightPayload [arrow]
Posted by "lidavidm (via GitHub)" <gi...@apache.org>.
lidavidm commented on PR #47115:
URL: https://github.com/apache/arrow/pull/47115#issuecomment-3198698555
Though, we should clarify what kind of stats Rusty actually wants
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] GH-45055: [C++][Flight] Update Flight Server RecordBatchStreamImpl to reuse ipc::RecordBatchWriter with custom IpcPayloadWriter instead of manually generating FlightPayload [arrow]
Posted by "lidavidm (via GitHub)" <gi...@apache.org>.
lidavidm commented on PR #47115:
URL: https://github.com/apache/arrow/pull/47115#issuecomment-3115370308
Thanks for tackling this Raúl!
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] GH-45055: [C++][Flight] Update Flight Server RecordBatchStreamImpl to reuse ipc::RecordBatchWriter with custom IpcPayloadWriter instead of manually generating FlightPayload [arrow]
Posted by "raulcd (via GitHub)" <gi...@apache.org>.
raulcd commented on PR #47115:
URL: https://github.com/apache/arrow/pull/47115#issuecomment-3197717756
that's interesting, I hacked something quick and dirty locally to validate dictionary deltas with flight just exposing the ` virtual arrow::ipc::ReadStats stats() const = 0;` to the `FlightStreamReader` and was able to validate deltas work (of course using the correct `IpcWriteOptions.emit_dictionary_deltas`:
```
pyarrow/tests/test_flight.py Stats: {'num_messages': 7, 'num_record_batches': 3, 'num_dictionary_batches': 3, 'num_dictionary_deltas': 2, 'num_replaced_dictionaries': 0}
```
@lidavidm what do you think about exposing some `ReadStats` on the `FlightStreamReader`, Should this be it's own FlightStats instead of `arrow::ipc::ReadStats`?
I recall there were some requests to have a method exposing the stats like, which we can use to do this:
- https://github.com/apache/arrow/issues/47171
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] GH-45055: [C++][Flight] Update Flight Server RecordBatchStreamImpl to reuse ipc::RecordBatchWriter with custom IpcPayloadWriter instead of manually generating FlightPayload [arrow]
Posted by "pitrou (via GitHub)" <gi...@apache.org>.
pitrou commented on PR #47115:
URL: https://github.com/apache/arrow/pull/47115#issuecomment-3197407867
> Should I add a new test for that?
If that's easy to do in Python then why not. The IPC reader exposes some useful `ReadStats`, but AFAICT the Flight client does not currently allows getting these stats. Perhaps it can be added?
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] GH-45055: [C++][Flight] Update Flight Server RecordBatchStreamImpl to reuse ipc::RecordBatchWriter with custom IpcPayloadWriter instead of manually generating FlightPayload [arrow]
Posted by "conbench-apache-arrow[bot] (via GitHub)" <gi...@apache.org>.
conbench-apache-arrow[bot] commented on PR #47115:
URL: https://github.com/apache/arrow/pull/47115#issuecomment-3138086343
After merging your PR, Conbench analyzed the 4 benchmarking runs that have been run so far on merge-commit 25f8f008061b978137a2c1e5bc934d07fae56e3e.
There weren't enough matching historic benchmark results to make a call on whether there were regressions.
The [full Conbench report](https://github.com/apache/arrow/runs/47076337452) has more details.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] GH-45055: [C++][Flight] Update Flight Server RecordBatchStreamImpl to reuse ipc::RecordBatchWriter with custom IpcPayloadWriter instead of manually generating FlightPayload [arrow]
Posted by "pitrou (via GitHub)" <gi...@apache.org>.
pitrou commented on PR #47115:
URL: https://github.com/apache/arrow/pull/47115#issuecomment-3196776330
For the record, does this also enable dictionary deltas or just dictionary replacement?
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] GH-45055: [C++][Flight] Update Flight Server RecordBatchStreamImpl to reuse ipc::RecordBatchWriter with custom IpcPayloadWriter instead of manually generating FlightPayload [arrow]
Posted by "raulcd (via GitHub)" <gi...@apache.org>.
raulcd commented on code in PR #47115:
URL: https://github.com/apache/arrow/pull/47115#discussion_r2224854056
##########
cpp/src/arrow/flight/server.cc:
##########
@@ -275,78 +277,137 @@ Status FlightServerBase::GetSchema(const ServerCallContext& context,
class RecordBatchStream::RecordBatchStreamImpl {
public:
- // Stages of the stream when producing payloads
- enum class Stage {
- NEW, // The stream has been created, but Next has not been called yet
- DICTIONARY, // Dictionaries have been collected, and are being sent
- RECORD_BATCH // Initial have been sent
- };
-
RecordBatchStreamImpl(const std::shared_ptr<RecordBatchReader>& reader,
const ipc::IpcWriteOptions& options)
- : reader_(reader), mapper_(*reader_->schema()), ipc_options_(options) {}
+ : reader_(reader), options_(options) {}
std::shared_ptr<Schema> schema() { return reader_->schema(); }
Status GetSchemaPayload(FlightPayload* payload) {
- return ipc::GetSchemaPayload(*reader_->schema(), ipc_options_, mapper_,
- &payload->ipc_message);
+ if (!writer_) {
+ // Create the IPC writer on first call
+ auto payload_writer =
+ std::make_unique<ServerRecordBatchPayloadWriter>(&payload_list_);
+ auto writer_result = ipc::internal::OpenRecordBatchWriter(
+ std::move(payload_writer), reader_->schema(), options_);
+
+ if (!writer_result.ok()) {
+ return writer_result.status();
+ }
+ writer_ = std::move(writer_result).ValueOrDie();
Review Comment:
Because I am still learning C++ :)
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] GH-45055: [C++][Flight] Update Flight Server RecordBatchStreamImpl to reuse ipc::RecordBatchWriter with custom IpcPayloadWriter instead of manually generating FlightPayload [arrow]
Posted by "raulcd (via GitHub)" <gi...@apache.org>.
raulcd commented on code in PR #47115:
URL: https://github.com/apache/arrow/pull/47115#discussion_r2225045071
##########
cpp/src/arrow/flight/server.cc:
##########
@@ -275,78 +277,137 @@ Status FlightServerBase::GetSchema(const ServerCallContext& context,
class RecordBatchStream::RecordBatchStreamImpl {
public:
- // Stages of the stream when producing payloads
- enum class Stage {
- NEW, // The stream has been created, but Next has not been called yet
- DICTIONARY, // Dictionaries have been collected, and are being sent
- RECORD_BATCH // Initial have been sent
- };
-
RecordBatchStreamImpl(const std::shared_ptr<RecordBatchReader>& reader,
const ipc::IpcWriteOptions& options)
- : reader_(reader), mapper_(*reader_->schema()), ipc_options_(options) {}
+ : reader_(reader), options_(options) {}
std::shared_ptr<Schema> schema() { return reader_->schema(); }
Status GetSchemaPayload(FlightPayload* payload) {
- return ipc::GetSchemaPayload(*reader_->schema(), ipc_options_, mapper_,
- &payload->ipc_message);
+ if (!writer_) {
+ // Create the IPC writer on first call
+ auto payload_writer =
+ std::make_unique<ServerRecordBatchPayloadWriter>(&payload_list_);
+ auto writer_result = ipc::internal::OpenRecordBatchWriter(
+ std::move(payload_writer), reader_->schema(), options_);
+
+ if (!writer_result.ok()) {
+ return writer_result.status();
+ }
+ writer_ = std::move(writer_result).ValueOrDie();
+ }
+
+ // Return the current payload (schema)
+ if (!payload_list_.empty()) {
+ *payload = std::move(payload_list_.front());
+ payload_list_.pop_front();
+ return Status::OK();
+ }
+ return Status::UnknownError("No schema payload generated");
}
Status Next(FlightPayload* payload) {
- if (stage_ == Stage::NEW) {
- RETURN_NOT_OK(reader_->ReadNext(¤t_batch_));
- if (!current_batch_) {
- // Signal that iteration is over
- payload->ipc_message.metadata = nullptr;
- return Status::OK();
+ // If we have previous payloads (dictionary messages or previous record batches)
+ // return them first before reading next record batch.
+ if (!payload_list_.empty()) {
+ *payload = std::move(payload_list_.front());
+ payload_list_.pop_front();
+ return Status::OK();
+ }
+
+ std::shared_ptr<RecordBatch> batch;
+ RETURN_NOT_OK(reader_->ReadNext(&batch));
+
+ if (!batch) {
+ // End of stream
+ if (writer_) {
+ RETURN_NOT_OK(writer_->Close());
}
- ARROW_ASSIGN_OR_RAISE(dictionaries_,
- ipc::CollectDictionaries(*current_batch_, mapper_));
- stage_ = Stage::DICTIONARY;
+ payload->ipc_message.metadata = nullptr;
+ return Status::OK();
}
- if (stage_ == Stage::DICTIONARY) {
- if (dictionary_index_ == static_cast<int>(dictionaries_.size())) {
- stage_ = Stage::RECORD_BATCH;
- return ipc::GetRecordBatchPayload(*current_batch_, ipc_options_,
- &payload->ipc_message);
- } else {
- return GetNextDictionary(payload);
+ // Check if writer is already initialized or if schema has changed.
+ // To recreate the writer.
Review Comment:
The test is `TestFlightSqlServer.TestCommandGetTablesWithIncludedSchemas` and failed with:
```c++
/home/raulcd/code/arrow/cpp/src/arrow/flight/sql/example/sqlite_server.cc:395: GetTables include_schema=1
/home/raulcd/code/arrow/cpp/src/arrow/flight/sql/example/sqlite_server.cc:409: GetTables: SELECT 'main' as catalog_name, null as schema_name, name as table_name, type as table_type FROM sqlite_master where 1=1 and table_name LIKE 'int%' order by table_name
Reading next payload from data stream: 1
/home/raulcd/code/arrow/cpp/src/arrow/flight/sql/server_test.cc:325: Failure
Failed
'_error_or_value62.status()' failed with Invalid: Tried to write record batch with different schema
/home/raulcd/code/arrow/cpp/src/arrow/flight/server.cc:339 writer_->WriteRecordBatch(*batch)
/home/raulcd/code/arrow/cpp/src/arrow/flight/server.cc:415 impl_->Next(&payload)
/home/raulcd/code/arrow/cpp/src/arrow/flight/transport_server.cc:291 data_stream->Next(). gRPC client debug context: {"created":"@1753261907.620461787","description":"Error received from peer ipv4:127.0.0.1:40331","file":"/home/raulcd/code/arrow/cpp/build/grpc_ep-prefix/src/grpc_ep/src/core/lib/surface/call.cc","file_line":952,"grpc_message":"Tried to write record batch with different schema\n/home/raulcd/code/arrow/cpp/src/arrow/flight/server.cc:339 writer_->WriteRecordBatch(*batch)\n/home/raulcd/code/arrow/cpp/src/arrow/flight/server.cc:415 impl_->Next(&payload)\n/home/raulcd/code/arrow/cpp/src/arrow/flight/transport_server.cc:291 data_stream->Next()","grpc_status":3}. Client context: OK
/home/raulcd/code/arrow/cpp/src/arrow/flight/client.cc:219 stream_->Finish(Status::OK())
/home/raulcd/code/arrow/cpp/src/arrow/flight/client.cc:256 Next()
/home/raulcd/code/arrow/cpp/src/arrow/flight/client.cc:116 ToRecordBatches(stop_token)
```
The schemas difference was `table_schema` from `binary` to `binary not null`:
```md
## Batch schema:
catalog_name: string
db_schema_name: string
table_name: string not null
table_type: string not null
table_schema: binary
## Reader Schema:
catalog_name: string
db_schema_name: string
table_name: string not null
table_type: string not null
table_schema: binary not null
```
I've found the issue and fixed here: https://github.com/apache/arrow/pull/47115/commits/d4305a8de40802db19bfaa27528bd94203219a56
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] GH-45055: [C++][Flight] Update Flight Server RecordBatchStreamImpl to reuse ipc::RecordBatchWriter with custom IpcPayloadWriter instead of manually generating FlightPayload [arrow]
Posted by "raulcd (via GitHub)" <gi...@apache.org>.
raulcd commented on PR #47115:
URL: https://github.com/apache/arrow/pull/47115#issuecomment-3197269948
Both as it uses the underlying `RecordBatchWriter` here:
https://github.com/apache/arrow/blob/71a7b5531fd6ef96ce52d04491656031eb023c09/cpp/src/arrow/ipc/writer.cc#L1272-L1340
It will generate the dictionary delta as we do on the Ipc RecordBatchWriter.
Should I add a new test for that?
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] GH-45055: [C++][Flight] Support dictionary replacement in DoGet [arrow]
Posted by "github-actions[bot] (via GitHub)" <gi...@apache.org>.
github-actions[bot] commented on PR #47115:
URL: https://github.com/apache/arrow/pull/47115#issuecomment-3074447325
:warning: GitHub issue #45055 **has been automatically assigned in GitHub** to PR creator.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] GH-45055: [C++][Flight] Update Flight Server RecordBatchStreamImpl to reuse ipc::RecordBatchWriter with custom IpcPayloadWriter instead of manually generating FlightPayload [arrow]
Posted by "lidavidm (via GitHub)" <gi...@apache.org>.
lidavidm commented on code in PR #47115:
URL: https://github.com/apache/arrow/pull/47115#discussion_r2224098026
##########
cpp/src/arrow/flight/server.cc:
##########
@@ -275,78 +277,137 @@ Status FlightServerBase::GetSchema(const ServerCallContext& context,
class RecordBatchStream::RecordBatchStreamImpl {
public:
- // Stages of the stream when producing payloads
- enum class Stage {
- NEW, // The stream has been created, but Next has not been called yet
- DICTIONARY, // Dictionaries have been collected, and are being sent
- RECORD_BATCH // Initial have been sent
- };
-
RecordBatchStreamImpl(const std::shared_ptr<RecordBatchReader>& reader,
const ipc::IpcWriteOptions& options)
- : reader_(reader), mapper_(*reader_->schema()), ipc_options_(options) {}
+ : reader_(reader), options_(options) {}
std::shared_ptr<Schema> schema() { return reader_->schema(); }
Status GetSchemaPayload(FlightPayload* payload) {
- return ipc::GetSchemaPayload(*reader_->schema(), ipc_options_, mapper_,
- &payload->ipc_message);
+ if (!writer_) {
+ // Create the IPC writer on first call
+ auto payload_writer =
+ std::make_unique<ServerRecordBatchPayloadWriter>(&payload_list_);
+ auto writer_result = ipc::internal::OpenRecordBatchWriter(
+ std::move(payload_writer), reader_->schema(), options_);
+
+ if (!writer_result.ok()) {
+ return writer_result.status();
+ }
+ writer_ = std::move(writer_result).ValueOrDie();
+ }
+
+ // Return the current payload (schema)
+ if (!payload_list_.empty()) {
+ *payload = std::move(payload_list_.front());
+ payload_list_.pop_front();
+ return Status::OK();
+ }
+ return Status::UnknownError("No schema payload generated");
}
Status Next(FlightPayload* payload) {
- if (stage_ == Stage::NEW) {
- RETURN_NOT_OK(reader_->ReadNext(¤t_batch_));
- if (!current_batch_) {
- // Signal that iteration is over
- payload->ipc_message.metadata = nullptr;
- return Status::OK();
+ // If we have previous payloads (dictionary messages or previous record batches)
+ // return them first before reading next record batch.
+ if (!payload_list_.empty()) {
+ *payload = std::move(payload_list_.front());
+ payload_list_.pop_front();
+ return Status::OK();
+ }
+
+ std::shared_ptr<RecordBatch> batch;
+ RETURN_NOT_OK(reader_->ReadNext(&batch));
+
+ if (!batch) {
+ // End of stream
+ if (writer_) {
+ RETURN_NOT_OK(writer_->Close());
}
- ARROW_ASSIGN_OR_RAISE(dictionaries_,
- ipc::CollectDictionaries(*current_batch_, mapper_));
- stage_ = Stage::DICTIONARY;
+ payload->ipc_message.metadata = nullptr;
+ return Status::OK();
}
- if (stage_ == Stage::DICTIONARY) {
- if (dictionary_index_ == static_cast<int>(dictionaries_.size())) {
- stage_ = Stage::RECORD_BATCH;
- return ipc::GetRecordBatchPayload(*current_batch_, ipc_options_,
- &payload->ipc_message);
- } else {
- return GetNextDictionary(payload);
+ // Check if writer is already initialized or if schema has changed.
+ // To recreate the writer.
+ if (!writer_ || !batch->schema()->Equals(*reader_->schema())) {
+ if (writer_) {
+ RETURN_NOT_OK(writer_->Close());
+ }
+
+ // Create new writer with the batch's schema
+ auto payload_writer =
+ std::make_unique<ServerRecordBatchPayloadWriter>(&payload_list_);
+ auto writer_result = ipc::internal::OpenRecordBatchWriter(
+ std::move(payload_writer), batch->schema(), options_);
+
+ if (!writer_result.ok()) {
+ return writer_result.status();
+ }
+ writer_ = std::move(writer_result).ValueOrDie();
+ if (!payload_list_.empty()) {
+ // Drop Schema message if it was generated.
+ // We want new Dictionary or RecordBatch messages only.
+ payload_list_.pop_front();
}
}
- RETURN_NOT_OK(reader_->ReadNext(¤t_batch_));
+ // One WriteRecordBatch call might generate multiple payloads, so we
+ // need to collect them in a list.
+ RETURN_NOT_OK(writer_->WriteRecordBatch(*batch));
- // TODO(ARROW-10787): Delta dictionaries
- if (!current_batch_) {
- // Signal that iteration is over
- payload->ipc_message.metadata = nullptr;
+ // Return the first generated payload after WriteRecordBatch. There must be
+ // at least one payload generated.
+ if (!payload_list_.empty()) {
+ *payload = std::move(payload_list_.front());
+ payload_list_.pop_front();
return Status::OK();
- } else {
- return ipc::GetRecordBatchPayload(*current_batch_, ipc_options_,
- &payload->ipc_message);
}
+
+ return Status::UnknownError("IPC writer didn't produce any payloads");
}
- Status Close() { return reader_->Close(); }
+ Status Close() {
+ if (writer_) {
+ RETURN_NOT_OK(writer_->Close());
+ }
+ return reader_->Close();
+ }
private:
- Status GetNextDictionary(FlightPayload* payload) {
- const auto& it = dictionaries_[dictionary_index_++];
- return ipc::GetDictionaryPayload(it.first, it.second, ipc_options_,
- &payload->ipc_message);
- }
+ // Simple payload writer that uses a list of generated payloads.
+ class ServerRecordBatchPayloadWriter : public ipc::internal::IpcPayloadWriter {
+ public:
+ explicit ServerRecordBatchPayloadWriter(std::list<FlightPayload>* payload_list)
+ : payload_list_(payload_list), first_payload_(true) {}
+
+ Status Start() override { return Status::OK(); }
+
+ Status WritePayload(const ipc::IpcPayload& ipc_payload) override {
+ FlightPayload payload;
+ payload.ipc_message = ipc_payload;
+
+ if (first_payload_) {
+ if (ipc_payload.type != ipc::MessageType::SCHEMA) {
+ return Status::Invalid("First IPC message should be schema");
+ }
+ first_payload_ = false;
+ }
Review Comment:
Do we really need to track first_payload_ and this, or just trust that the IPC writer does its job?
##########
cpp/src/arrow/flight/server.cc:
##########
@@ -275,78 +277,137 @@ Status FlightServerBase::GetSchema(const ServerCallContext& context,
class RecordBatchStream::RecordBatchStreamImpl {
public:
- // Stages of the stream when producing payloads
- enum class Stage {
- NEW, // The stream has been created, but Next has not been called yet
- DICTIONARY, // Dictionaries have been collected, and are being sent
- RECORD_BATCH // Initial have been sent
- };
-
RecordBatchStreamImpl(const std::shared_ptr<RecordBatchReader>& reader,
const ipc::IpcWriteOptions& options)
- : reader_(reader), mapper_(*reader_->schema()), ipc_options_(options) {}
+ : reader_(reader), options_(options) {}
std::shared_ptr<Schema> schema() { return reader_->schema(); }
Status GetSchemaPayload(FlightPayload* payload) {
- return ipc::GetSchemaPayload(*reader_->schema(), ipc_options_, mapper_,
- &payload->ipc_message);
+ if (!writer_) {
+ // Create the IPC writer on first call
+ auto payload_writer =
+ std::make_unique<ServerRecordBatchPayloadWriter>(&payload_list_);
+ auto writer_result = ipc::internal::OpenRecordBatchWriter(
+ std::move(payload_writer), reader_->schema(), options_);
+
+ if (!writer_result.ok()) {
+ return writer_result.status();
+ }
+ writer_ = std::move(writer_result).ValueOrDie();
+ }
+
+ // Return the current payload (schema)
+ if (!payload_list_.empty()) {
+ *payload = std::move(payload_list_.front());
+ payload_list_.pop_front();
+ return Status::OK();
+ }
+ return Status::UnknownError("No schema payload generated");
}
Status Next(FlightPayload* payload) {
- if (stage_ == Stage::NEW) {
- RETURN_NOT_OK(reader_->ReadNext(¤t_batch_));
- if (!current_batch_) {
- // Signal that iteration is over
- payload->ipc_message.metadata = nullptr;
- return Status::OK();
+ // If we have previous payloads (dictionary messages or previous record batches)
+ // return them first before reading next record batch.
+ if (!payload_list_.empty()) {
+ *payload = std::move(payload_list_.front());
+ payload_list_.pop_front();
+ return Status::OK();
+ }
Review Comment:
You could do something like
```
if (payload_list_.empty()) {
// Invoke the writer
}
if (payload_list_.empty()) {
// Error
}
// Return first payload
##########
cpp/src/arrow/flight/server.cc:
##########
@@ -275,78 +277,139 @@ Status FlightServerBase::GetSchema(const ServerCallContext& context,
class RecordBatchStream::RecordBatchStreamImpl {
public:
- // Stages of the stream when producing payloads
- enum class Stage {
- NEW, // The stream has been created, but Next has not been called yet
- DICTIONARY, // Dictionaries have been collected, and are being sent
- RECORD_BATCH // Initial have been sent
- };
-
RecordBatchStreamImpl(const std::shared_ptr<RecordBatchReader>& reader,
const ipc::IpcWriteOptions& options)
- : reader_(reader), mapper_(*reader_->schema()), ipc_options_(options) {}
+ : reader_(reader), options_(options) {}
std::shared_ptr<Schema> schema() { return reader_->schema(); }
Status GetSchemaPayload(FlightPayload* payload) {
- return ipc::GetSchemaPayload(*reader_->schema(), ipc_options_, mapper_,
- &payload->ipc_message);
+ if (!writer_) {
+ // Create the IPC writer on first call
+ auto payload_writer =
+ std::make_unique<ServerRecordBatchPayloadWriter>(&payload_list_);
+ auto writer_result = ipc::internal::OpenRecordBatchWriter(
+ std::move(payload_writer), reader_->schema(), options_);
+
+ if (!writer_result.ok()) {
+ return writer_result.status();
+ }
+ writer_ = std::move(writer_result).ValueOrDie();
+ }
+
+ // Return the current payload (schema)
+ if (!payload_list_.empty()) {
+ *payload = std::move(payload_list_.front());
+ payload_list_.pop_front();
+ return Status::OK();
+ }
+ return Status::UnknownError("No schema payload generated");
}
Status Next(FlightPayload* payload) {
- if (stage_ == Stage::NEW) {
- RETURN_NOT_OK(reader_->ReadNext(¤t_batch_));
- if (!current_batch_) {
- // Signal that iteration is over
- payload->ipc_message.metadata = nullptr;
- return Status::OK();
+ // If we have previous payloads (dictionary messages or previous record batches)
+ // return them first before reading next record batch.
+ if (!payload_list_.empty()) {
+ *payload = std::move(payload_list_.front());
+ payload_list_.pop_front();
+ return Status::OK();
+ }
+
+ std::shared_ptr<RecordBatch> batch;
+ RETURN_NOT_OK(reader_->ReadNext(&batch));
+
+ if (!batch) {
+ // End of stream
+ if (writer_) {
+ RETURN_NOT_OK(writer_->Close());
}
- ARROW_ASSIGN_OR_RAISE(dictionaries_,
- ipc::CollectDictionaries(*current_batch_, mapper_));
- stage_ = Stage::DICTIONARY;
+ payload->ipc_message.metadata = nullptr;
+ return Status::OK();
}
- if (stage_ == Stage::DICTIONARY) {
- if (dictionary_index_ == static_cast<int>(dictionaries_.size())) {
- stage_ = Stage::RECORD_BATCH;
- return ipc::GetRecordBatchPayload(*current_batch_, ipc_options_,
- &payload->ipc_message);
- } else {
- return GetNextDictionary(payload);
+ // Check if writer is already initialized or if schema has changed.
+ // To recreate the writer.
+ // TODO: Investigate why schema changed is needed for a flight-sql specific test.
Review Comment:
Which test? I don't think this is expected/they probably are buggy, in that we never validated things before
##########
cpp/src/arrow/flight/server.cc:
##########
@@ -275,78 +277,137 @@ Status FlightServerBase::GetSchema(const ServerCallContext& context,
class RecordBatchStream::RecordBatchStreamImpl {
public:
- // Stages of the stream when producing payloads
- enum class Stage {
- NEW, // The stream has been created, but Next has not been called yet
- DICTIONARY, // Dictionaries have been collected, and are being sent
- RECORD_BATCH // Initial have been sent
- };
-
RecordBatchStreamImpl(const std::shared_ptr<RecordBatchReader>& reader,
const ipc::IpcWriteOptions& options)
- : reader_(reader), mapper_(*reader_->schema()), ipc_options_(options) {}
+ : reader_(reader), options_(options) {}
std::shared_ptr<Schema> schema() { return reader_->schema(); }
Status GetSchemaPayload(FlightPayload* payload) {
- return ipc::GetSchemaPayload(*reader_->schema(), ipc_options_, mapper_,
- &payload->ipc_message);
+ if (!writer_) {
+ // Create the IPC writer on first call
+ auto payload_writer =
+ std::make_unique<ServerRecordBatchPayloadWriter>(&payload_list_);
+ auto writer_result = ipc::internal::OpenRecordBatchWriter(
+ std::move(payload_writer), reader_->schema(), options_);
+
+ if (!writer_result.ok()) {
+ return writer_result.status();
+ }
+ writer_ = std::move(writer_result).ValueOrDie();
Review Comment:
Why not ARROW_ASSIGN_OR_RAISE?
##########
cpp/src/arrow/flight/server.cc:
##########
@@ -275,78 +277,137 @@ Status FlightServerBase::GetSchema(const ServerCallContext& context,
class RecordBatchStream::RecordBatchStreamImpl {
public:
- // Stages of the stream when producing payloads
- enum class Stage {
- NEW, // The stream has been created, but Next has not been called yet
- DICTIONARY, // Dictionaries have been collected, and are being sent
- RECORD_BATCH // Initial have been sent
- };
-
RecordBatchStreamImpl(const std::shared_ptr<RecordBatchReader>& reader,
const ipc::IpcWriteOptions& options)
- : reader_(reader), mapper_(*reader_->schema()), ipc_options_(options) {}
+ : reader_(reader), options_(options) {}
std::shared_ptr<Schema> schema() { return reader_->schema(); }
Status GetSchemaPayload(FlightPayload* payload) {
- return ipc::GetSchemaPayload(*reader_->schema(), ipc_options_, mapper_,
- &payload->ipc_message);
+ if (!writer_) {
+ // Create the IPC writer on first call
+ auto payload_writer =
+ std::make_unique<ServerRecordBatchPayloadWriter>(&payload_list_);
+ auto writer_result = ipc::internal::OpenRecordBatchWriter(
+ std::move(payload_writer), reader_->schema(), options_);
+
+ if (!writer_result.ok()) {
+ return writer_result.status();
+ }
+ writer_ = std::move(writer_result).ValueOrDie();
+ }
+
+ // Return the current payload (schema)
+ if (!payload_list_.empty()) {
+ *payload = std::move(payload_list_.front());
+ payload_list_.pop_front();
+ return Status::OK();
+ }
+ return Status::UnknownError("No schema payload generated");
}
Status Next(FlightPayload* payload) {
- if (stage_ == Stage::NEW) {
- RETURN_NOT_OK(reader_->ReadNext(¤t_batch_));
- if (!current_batch_) {
- // Signal that iteration is over
- payload->ipc_message.metadata = nullptr;
- return Status::OK();
+ // If we have previous payloads (dictionary messages or previous record batches)
+ // return them first before reading next record batch.
+ if (!payload_list_.empty()) {
+ *payload = std::move(payload_list_.front());
+ payload_list_.pop_front();
+ return Status::OK();
+ }
+
+ std::shared_ptr<RecordBatch> batch;
+ RETURN_NOT_OK(reader_->ReadNext(&batch));
+
+ if (!batch) {
+ // End of stream
+ if (writer_) {
+ RETURN_NOT_OK(writer_->Close());
}
- ARROW_ASSIGN_OR_RAISE(dictionaries_,
- ipc::CollectDictionaries(*current_batch_, mapper_));
- stage_ = Stage::DICTIONARY;
+ payload->ipc_message.metadata = nullptr;
+ return Status::OK();
}
- if (stage_ == Stage::DICTIONARY) {
- if (dictionary_index_ == static_cast<int>(dictionaries_.size())) {
- stage_ = Stage::RECORD_BATCH;
- return ipc::GetRecordBatchPayload(*current_batch_, ipc_options_,
- &payload->ipc_message);
- } else {
- return GetNextDictionary(payload);
+ // Check if writer is already initialized or if schema has changed.
+ // To recreate the writer.
Review Comment:
I'm not actually sure what that test is about...we don't support changing schemas on the same stream (this was discussed on the ML a while back and discarded for lack of anyone to push it forward with all the bikeshedding)
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] GH-45055: [C++][Flight] Update Flight Server RecordBatchStreamImpl to reuse ipc::RecordBatchWriter with custom IpcPayloadWriter instead of manually generating FlightPayload [arrow]
Posted by "lidavidm (via GitHub)" <gi...@apache.org>.
lidavidm commented on PR #47115:
URL: https://github.com/apache/arrow/pull/47115#issuecomment-3198697902
I think we can just use the regular stats
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] GH-45055: [C++][Flight] Update Flight Server RecordBatchStreamImpl to reuse ipc::RecordBatchWriter with custom IpcPayloadWriter instead of manually generating FlightPayload [arrow]
Posted by "raulcd (via GitHub)" <gi...@apache.org>.
raulcd commented on PR #47115:
URL: https://github.com/apache/arrow/pull/47115#issuecomment-3096983775
@lidavidm I've tried to play some tricks to reuse the `ipc::RecordBatchWriter` with a new IpcPayloadWriter, similar to what we do on the client. This is the first time I am doing anything with Flight so I might have completely missed something, please let me know if that's the case :)
I found a couple of special cases from our tests that surprised me because they diverge from my expectations over IPC.
There's a flight sql test that seems to be changing the schema mid stream, is this expected? Can we return FlightPayloads that doesn't match the returned schema?
I am not adding more tests because with the change on the Python test I am already covering the original issue for the replacement dict on DoGet and with the previous tests we cover possible regressions. Let me know if you would like other test cases to be tackled.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] GH-45055: [C++][Flight] Update Flight Server RecordBatchStreamImpl to reuse ipc::RecordBatchWriter with custom IpcPayloadWriter instead of manually generating FlightPayload [arrow]
Posted by "raulcd (via GitHub)" <gi...@apache.org>.
raulcd commented on code in PR #47115:
URL: https://github.com/apache/arrow/pull/47115#discussion_r2219263167
##########
cpp/src/arrow/flight/server.cc:
##########
@@ -275,78 +277,139 @@ Status FlightServerBase::GetSchema(const ServerCallContext& context,
class RecordBatchStream::RecordBatchStreamImpl {
public:
- // Stages of the stream when producing payloads
- enum class Stage {
- NEW, // The stream has been created, but Next has not been called yet
- DICTIONARY, // Dictionaries have been collected, and are being sent
- RECORD_BATCH // Initial have been sent
- };
-
RecordBatchStreamImpl(const std::shared_ptr<RecordBatchReader>& reader,
const ipc::IpcWriteOptions& options)
- : reader_(reader), mapper_(*reader_->schema()), ipc_options_(options) {}
+ : reader_(reader), options_(options) {}
std::shared_ptr<Schema> schema() { return reader_->schema(); }
Status GetSchemaPayload(FlightPayload* payload) {
- return ipc::GetSchemaPayload(*reader_->schema(), ipc_options_, mapper_,
- &payload->ipc_message);
+ if (!writer_) {
+ // Create the IPC writer on first call
+ auto payload_writer =
+ std::make_unique<ServerRecordBatchPayloadWriter>(&payload_list_);
+ auto writer_result = ipc::internal::OpenRecordBatchWriter(
+ std::move(payload_writer), reader_->schema(), options_);
+
+ if (!writer_result.ok()) {
+ return writer_result.status();
+ }
+ writer_ = std::move(writer_result).ValueOrDie();
+ }
+
+ // Return the current payload (schema)
+ if (!payload_list_.empty()) {
+ *payload = std::move(payload_list_.front());
+ payload_list_.pop_front();
+ return Status::OK();
+ }
+ return Status::UnknownError("No schema payload generated");
}
Status Next(FlightPayload* payload) {
- if (stage_ == Stage::NEW) {
- RETURN_NOT_OK(reader_->ReadNext(¤t_batch_));
- if (!current_batch_) {
- // Signal that iteration is over
- payload->ipc_message.metadata = nullptr;
- return Status::OK();
+ // If we have previous payloads (dictionary messages or previous record batches)
+ // return them first before reading next record batch.
+ if (!payload_list_.empty()) {
+ *payload = std::move(payload_list_.front());
+ payload_list_.pop_front();
+ return Status::OK();
+ }
+
+ std::shared_ptr<RecordBatch> batch;
+ RETURN_NOT_OK(reader_->ReadNext(&batch));
+
+ if (!batch) {
+ // End of stream
+ if (writer_) {
+ RETURN_NOT_OK(writer_->Close());
}
- ARROW_ASSIGN_OR_RAISE(dictionaries_,
- ipc::CollectDictionaries(*current_batch_, mapper_));
- stage_ = Stage::DICTIONARY;
+ payload->ipc_message.metadata = nullptr;
+ return Status::OK();
}
- if (stage_ == Stage::DICTIONARY) {
- if (dictionary_index_ == static_cast<int>(dictionaries_.size())) {
- stage_ = Stage::RECORD_BATCH;
- return ipc::GetRecordBatchPayload(*current_batch_, ipc_options_,
- &payload->ipc_message);
- } else {
- return GetNextDictionary(payload);
+ // Check if writer is already initialized or if schema has changed.
+ // To recreate the writer.
+ // TODO: Investigate why schema changed is needed for a flight-sql specific test.
Review Comment:
Is a change of schema with an already generated writer expected? This was happening on a flight-sql test.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] GH-45055: [C++][Flight] Update Flight Server RecordBatchStreamImpl to reuse ipc::RecordBatchWriter with custom IpcPayloadWriter instead of manually generating FlightPayload [arrow]
Posted by "raulcd (via GitHub)" <gi...@apache.org>.
raulcd commented on PR #47115:
URL: https://github.com/apache/arrow/pull/47115#issuecomment-3112947087
Appveyor is timing out cloning the repository for the last builds (not only on this PR). It is not related to this PR and was successful on previous commits so CI is green.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] GH-45055: [C++][Flight] Update Flight Server RecordBatchStreamImpl to reuse ipc::RecordBatchWriter with custom IpcPayloadWriter instead of manually generating FlightPayload [arrow]
Posted by "raulcd (via GitHub)" <gi...@apache.org>.
raulcd commented on PR #47115:
URL: https://github.com/apache/arrow/pull/47115#issuecomment-3108674838
@lidavidm I think I've taken care of the initial comments! Thanks for those, I am learning quite a little bit with this PR so it was a fun one to work on!
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] GH-45055: [C++][Flight] Update Flight Server RecordBatchStreamImpl to reuse ipc::RecordBatchWriter with custom IpcPayloadWriter instead of manually generating FlightPayload [arrow]
Posted by "raulcd (via GitHub)" <gi...@apache.org>.
raulcd commented on code in PR #47115:
URL: https://github.com/apache/arrow/pull/47115#discussion_r2224847732
##########
cpp/src/arrow/flight/server.cc:
##########
@@ -275,78 +277,137 @@ Status FlightServerBase::GetSchema(const ServerCallContext& context,
class RecordBatchStream::RecordBatchStreamImpl {
public:
- // Stages of the stream when producing payloads
- enum class Stage {
- NEW, // The stream has been created, but Next has not been called yet
- DICTIONARY, // Dictionaries have been collected, and are being sent
- RECORD_BATCH // Initial have been sent
- };
-
RecordBatchStreamImpl(const std::shared_ptr<RecordBatchReader>& reader,
const ipc::IpcWriteOptions& options)
- : reader_(reader), mapper_(*reader_->schema()), ipc_options_(options) {}
+ : reader_(reader), options_(options) {}
std::shared_ptr<Schema> schema() { return reader_->schema(); }
Status GetSchemaPayload(FlightPayload* payload) {
- return ipc::GetSchemaPayload(*reader_->schema(), ipc_options_, mapper_,
- &payload->ipc_message);
+ if (!writer_) {
+ // Create the IPC writer on first call
+ auto payload_writer =
+ std::make_unique<ServerRecordBatchPayloadWriter>(&payload_list_);
+ auto writer_result = ipc::internal::OpenRecordBatchWriter(
+ std::move(payload_writer), reader_->schema(), options_);
+
+ if (!writer_result.ok()) {
+ return writer_result.status();
+ }
+ writer_ = std::move(writer_result).ValueOrDie();
+ }
+
+ // Return the current payload (schema)
+ if (!payload_list_.empty()) {
+ *payload = std::move(payload_list_.front());
+ payload_list_.pop_front();
+ return Status::OK();
+ }
+ return Status::UnknownError("No schema payload generated");
}
Status Next(FlightPayload* payload) {
- if (stage_ == Stage::NEW) {
- RETURN_NOT_OK(reader_->ReadNext(¤t_batch_));
- if (!current_batch_) {
- // Signal that iteration is over
- payload->ipc_message.metadata = nullptr;
- return Status::OK();
+ // If we have previous payloads (dictionary messages or previous record batches)
+ // return them first before reading next record batch.
+ if (!payload_list_.empty()) {
+ *payload = std::move(payload_list_.front());
+ payload_list_.pop_front();
+ return Status::OK();
+ }
+
+ std::shared_ptr<RecordBatch> batch;
+ RETURN_NOT_OK(reader_->ReadNext(&batch));
+
+ if (!batch) {
+ // End of stream
+ if (writer_) {
+ RETURN_NOT_OK(writer_->Close());
}
- ARROW_ASSIGN_OR_RAISE(dictionaries_,
- ipc::CollectDictionaries(*current_batch_, mapper_));
- stage_ = Stage::DICTIONARY;
+ payload->ipc_message.metadata = nullptr;
+ return Status::OK();
}
- if (stage_ == Stage::DICTIONARY) {
- if (dictionary_index_ == static_cast<int>(dictionaries_.size())) {
- stage_ = Stage::RECORD_BATCH;
- return ipc::GetRecordBatchPayload(*current_batch_, ipc_options_,
- &payload->ipc_message);
- } else {
- return GetNextDictionary(payload);
+ // Check if writer is already initialized or if schema has changed.
+ // To recreate the writer.
+ if (!writer_ || !batch->schema()->Equals(*reader_->schema())) {
+ if (writer_) {
+ RETURN_NOT_OK(writer_->Close());
+ }
+
+ // Create new writer with the batch's schema
+ auto payload_writer =
+ std::make_unique<ServerRecordBatchPayloadWriter>(&payload_list_);
+ auto writer_result = ipc::internal::OpenRecordBatchWriter(
+ std::move(payload_writer), batch->schema(), options_);
+
+ if (!writer_result.ok()) {
+ return writer_result.status();
+ }
+ writer_ = std::move(writer_result).ValueOrDie();
+ if (!payload_list_.empty()) {
+ // Drop Schema message if it was generated.
+ // We want new Dictionary or RecordBatch messages only.
+ payload_list_.pop_front();
}
}
- RETURN_NOT_OK(reader_->ReadNext(¤t_batch_));
+ // One WriteRecordBatch call might generate multiple payloads, so we
+ // need to collect them in a list.
+ RETURN_NOT_OK(writer_->WriteRecordBatch(*batch));
- // TODO(ARROW-10787): Delta dictionaries
- if (!current_batch_) {
- // Signal that iteration is over
- payload->ipc_message.metadata = nullptr;
+ // Return the first generated payload after WriteRecordBatch. There must be
+ // at least one payload generated.
+ if (!payload_list_.empty()) {
+ *payload = std::move(payload_list_.front());
+ payload_list_.pop_front();
return Status::OK();
- } else {
- return ipc::GetRecordBatchPayload(*current_batch_, ipc_options_,
- &payload->ipc_message);
}
+
+ return Status::UnknownError("IPC writer didn't produce any payloads");
}
- Status Close() { return reader_->Close(); }
+ Status Close() {
+ if (writer_) {
+ RETURN_NOT_OK(writer_->Close());
+ }
+ return reader_->Close();
+ }
private:
- Status GetNextDictionary(FlightPayload* payload) {
- const auto& it = dictionaries_[dictionary_index_++];
- return ipc::GetDictionaryPayload(it.first, it.second, ipc_options_,
- &payload->ipc_message);
- }
+ // Simple payload writer that uses a list of generated payloads.
+ class ServerRecordBatchPayloadWriter : public ipc::internal::IpcPayloadWriter {
+ public:
+ explicit ServerRecordBatchPayloadWriter(std::list<FlightPayload>* payload_list)
+ : payload_list_(payload_list), first_payload_(true) {}
+
+ Status Start() override { return Status::OK(); }
+
+ Status WritePayload(const ipc::IpcPayload& ipc_payload) override {
+ FlightPayload payload;
+ payload.ipc_message = ipc_payload;
+
+ if (first_payload_) {
+ if (ipc_payload.type != ipc::MessageType::SCHEMA) {
+ return Status::Invalid("First IPC message should be schema");
+ }
+ first_payload_ = false;
+ }
Review Comment:
Probably not, the IPC writer should guarantee that.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] GH-45055: [C++][Flight] Update Flight Server RecordBatchStreamImpl to reuse ipc::RecordBatchWriter with custom IpcPayloadWriter instead of manually generating FlightPayload [arrow]
Posted by "lidavidm (via GitHub)" <gi...@apache.org>.
lidavidm merged PR #47115:
URL: https://github.com/apache/arrow/pull/47115
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org