| 423 | AssertBatchesEqual(*new_batch, *batch2); |
| 424 | |
| 425 | ASSERT_OK_AND_ASSIGN(auto new_batch2, batch.AddColumn(1, "f3", array3)); |
| 426 | AssertBatchesEqual(*new_batch2, *new_batch); |
| 427 | |
| 428 | ASSERT_TRUE(new_batch2->schema()->field(1)->nullable()); |
| 429 | } |
| 430 | |
| 431 | TEST_F(TestRecordBatch, SetColumn) { |
| 432 | const int length = 10; |
| 433 | |
| 434 | auto field1 = field("f1", int32()); |
| 435 | auto field2 = field("f2", uint8()); |
| 436 | auto field3 = field("f3", int16()); |
| 437 | |
| 438 | auto schema1 = ::arrow::schema({field1, field2}); |
| 439 | auto schema2 = ::arrow::schema({field1, field3}); |
| 440 | auto schema3 = ::arrow::schema({field3, field2}); |
| 441 | |
| 442 | random::RandomArrayGenerator gen(42); |
| 443 | |
| 444 | auto array1 = gen.ArrayOf(int32(), length); |
| 445 | auto array2 = gen.ArrayOf(uint8(), length); |
| 446 | auto array3 = gen.ArrayOf(int16(), length); |
| 447 | |
| 448 | auto batch1 = RecordBatch::Make(schema1, length, {array1, array2}); |
| 449 | auto batch2 = RecordBatch::Make(schema2, length, {array1, array3}); |
| 450 | auto batch3 = RecordBatch::Make(schema3, length, {array3, array2}); |
| 451 | |
| 452 | const RecordBatch& batch = *batch1; |
| 453 | |
| 454 | // Negative tests with invalid index |
| 455 | ASSERT_RAISES(Invalid, batch.SetColumn(5, field1, array1)); |
| 456 | ASSERT_RAISES(Invalid, batch.SetColumn(-1, field1, array1)); |
| 457 | |
| 458 | // Negative test with wrong length |
| 459 | auto longer_col = gen.ArrayOf(int32(), length + 1); |
| 460 | ASSERT_RAISES(Invalid, batch.SetColumn(0, field1, longer_col)); |
| 461 | |
| 462 | // Negative test with mismatch type |
| 463 | ASSERT_RAISES(TypeError, batch.SetColumn(0, field1, array2)); |
| 464 | |
| 465 | ASSERT_OK_AND_ASSIGN(auto new_batch, batch.SetColumn(1, field3, array3)); |
| 466 | AssertBatchesEqual(*new_batch, *batch2); |
| 467 | |
| 468 | ASSERT_OK_AND_ASSIGN(new_batch, batch.SetColumn(0, field3, array3)); |
| 469 | AssertBatchesEqual(*new_batch, *batch3); |
| 470 | } |
| 471 | |
| 472 | TEST_F(TestRecordBatch, RemoveColumn) { |
| 473 | const int length = 10; |
| 474 | |
| 475 | auto field1 = field("f1", int32()); |
| 476 | auto field2 = field("f2", uint8()); |
| 477 | auto field3 = field("f3", int16()); |
| 478 | |
| 479 | auto schema1 = ::arrow::schema({field1, field2, field3}); |
| 480 | auto schema2 = ::arrow::schema({field2, field3}); |
| 481 | auto schema3 = ::arrow::schema({field1, field3}); |
| 482 | auto schema4 = ::arrow::schema({field1, field2}); |
no test coverage detected