| 4351 | |
| 4352 | template <typename BatchFactory> |
| 4353 | void TestWithBatchFactory(BatchFactory&& factory) { |
| 4354 | std::shared_ptr<Device> device = std::make_shared<MyDevice>(1); |
| 4355 | auto mm = device->default_memory_manager(); |
| 4356 | |
| 4357 | std::shared_ptr<RecordBatch> batch; |
| 4358 | struct ArrowDeviceArray c_array {}; |
| 4359 | struct ArrowSchema c_schema {}; |
| 4360 | ArrayExportGuard array_guard(&c_array.array); |
| 4361 | SchemaExportGuard schema_guard(&c_schema); |
| 4362 | |
| 4363 | auto orig_bytes = pool_->bytes_allocated(); |
| 4364 | ASSERT_OK_AND_ASSIGN(batch, ToResult(factory())); |
| 4365 | ASSERT_OK(ExportSchema(*batch->schema(), &c_schema)); |
| 4366 | ASSERT_OK_AND_ASSIGN(auto sync, mm->MakeDeviceSyncEvent()); |
| 4367 | ASSERT_OK(ExportDeviceRecordBatch(*batch, sync, &c_array)); |
| 4368 | |
| 4369 | auto new_bytes = pool_->bytes_allocated(); |
| 4370 | batch.reset(); |
| 4371 | ASSERT_EQ(pool_->bytes_allocated(), new_bytes); |
| 4372 | ASSERT_OK_AND_ASSIGN(batch, |
| 4373 | ImportDeviceRecordBatch(&c_array, &c_schema, DeviceMapper)); |
| 4374 | ASSERT_OK(batch->ValidateFull()); |
| 4375 | ASSERT_TRUE(ArrowSchemaIsReleased(&c_schema)); |
| 4376 | ASSERT_TRUE(ArrowArrayIsReleased(&c_array.array)); |
| 4377 | |
| 4378 | // Re-export and re-import, now both at once |
| 4379 | ASSERT_OK(ExportDeviceRecordBatch(*batch, sync, &c_array, &c_schema)); |
| 4380 | batch.reset(); |
| 4381 | ASSERT_OK_AND_ASSIGN(batch, |
| 4382 | ImportDeviceRecordBatch(&c_array, &c_schema, DeviceMapper)); |
| 4383 | ASSERT_OK(batch->ValidateFull()); |
| 4384 | ASSERT_TRUE(ArrowSchemaIsReleased(&c_schema)); |
| 4385 | ASSERT_TRUE(ArrowArrayIsReleased(&c_array.array)); |
| 4386 | |
| 4387 | // Check value of imported record batch |
| 4388 | { |
| 4389 | std::shared_ptr<RecordBatch> expected; |
| 4390 | ASSERT_OK_AND_ASSIGN(expected, ToResult(factory())); |
| 4391 | AssertSchemaEqual(*expected->schema(), *batch->schema(), /*check_metadata=*/true); |
| 4392 | AssertBatchesEqual(*expected, *batch); |
| 4393 | } |
| 4394 | batch.reset(); |
| 4395 | ASSERT_EQ(pool_->bytes_allocated(), orig_bytes); |
| 4396 | } |
| 4397 | |
| 4398 | void TestWithJSON(const std::shared_ptr<MemoryManager>& mm, |
| 4399 | std::shared_ptr<DataType> type, const char* json) { |
nothing calls this directly
no test coverage detected