pitrou commented on a change in pull request #12628:
URL: https://github.com/apache/arrow/pull/12628#discussion_r829177708
##########
File path: cpp/src/arrow/flight/test_definitions.cc
##########
@@ -1140,17 +1271,34 @@ void CudaDataTest::TestDoGet() {
FlightCallOptions options;
options.memory_manager = impl_->device->default_memory_manager();
+ const RecordBatchVector& batches =
+ reinterpret_cast<CudaTestServer*>(server_.get())->batches();
+
Ticket ticket{""};
std::unique_ptr<FlightStreamReader> stream;
ASSERT_OK(client_->DoGet(options, ticket, &stream));
- std::shared_ptr<Table> table;
- ASSERT_OK(stream->ReadAll(&table));
- for (const auto& column : table->columns()) {
- for (const auto& chunk : column->chunks()) {
- ASSERT_OK(CheckBuffersOnDevice(*chunk, *impl_->device));
+ size_t idx = 0;
+ while (true) {
+ FlightStreamChunk chunk;
+ ASSERT_OK(stream->Next(&chunk));
+ if (!chunk.data) break;
+
+ for (const auto& column : chunk.data->columns()) {
+ ASSERT_OK(CheckBuffersOnDevice(*column, *impl_->device));
Review comment:
Should we have a `CheckBuffersOnDevice` overload for record batches?
##########
File path: cpp/src/arrow/flight/test_definitions.cc
##########
@@ -518,22 +589,49 @@ void DataTest::TestIssue5095() {
//------------------------------------------------------------
// Specific tests for DoPut
+static constexpr char kExpectedMetadata[] = "foo bar";
+
class DoPutTestServer : public FlightServerBase {
public:
Status DoPut(const ServerCallContext& context,
std::unique_ptr<FlightMessageReader> reader,
std::unique_ptr<FlightMetadataWriter> writer) override {
descriptor_ = reader->descriptor();
+
+ if (descriptor_.type == FlightDescriptor::DescriptorType::CMD) {
+ if (descriptor_.cmd == "TestUndrained") {
+ // Don't read all the messages
+ return Status::OK();
+ }
+ }
+
int counter = 0;
+ FlightStreamChunk chunk;
while (true) {
- FlightStreamChunk chunk;
RETURN_NOT_OK(reader->Next(&chunk));
if (!chunk.data) break;
+ if (counter % 2 == 1) {
Review comment:
What happens if `counter % 2 == 0`? Should we assert that there's no app
metadata?
##########
File path: cpp/src/arrow/flight/test_definitions.cc
##########
@@ -1140,17 +1271,34 @@ void CudaDataTest::TestDoGet() {
FlightCallOptions options;
options.memory_manager = impl_->device->default_memory_manager();
+ const RecordBatchVector& batches =
+ reinterpret_cast<CudaTestServer*>(server_.get())->batches();
Review comment:
Should we use `checked_cast` here?
##########
File path: cpp/src/arrow/flight/test_definitions.cc
##########
@@ -1175,6 +1323,23 @@ void CudaDataTest::TestDoPut() {
ASSERT_OK(writer->WriteRecordBatch(*cuda_batch));
}
ASSERT_OK(writer->Close());
+ ASSERT_OK(impl_->context->Synchronize());
+
+ const RecordBatchVector& written =
+ reinterpret_cast<CudaTestServer*>(server_.get())->batches();
Review comment:
`checked_cast` as well?
--
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: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
