| 60 | }; |
| 61 | |
| 62 | TEST_F(SimpleWriteNodeTest, CustomNullability) { |
| 63 | // Create an input table with a nullable and a non-nullable type |
| 64 | ExecBatch batch = gen::Gen({gen::Step()})->FailOnError()->ExecBatch(/*num_rows=*/1); |
| 65 | std::shared_ptr<Schema> test_schema = |
| 66 | schema({field("nullable_i32", uint32(), /*nullable=*/true), |
| 67 | field("non_nullable_i32", uint32(), /*nullable=*/false)}); |
| 68 | std::shared_ptr<RecordBatch> record_batch = |
| 69 | RecordBatch::Make(test_schema, /*num_rows=*/1, |
| 70 | {batch.values[0].make_array(), batch.values[0].make_array()}); |
| 71 | ASSERT_OK_AND_ASSIGN(std::shared_ptr<Table> table, |
| 72 | Table::FromRecordBatches({std::move(record_batch)})); |
| 73 | |
| 74 | ASSERT_TRUE(table->field(0)->nullable()); |
| 75 | ASSERT_FALSE(table->field(1)->nullable()); |
| 76 | |
| 77 | dataset::WriteNodeOptions write_options(fs_write_options_); |
| 78 | write_options.custom_schema = test_schema; |
| 79 | |
| 80 | // Write the data to disk (these plans use a project because it destroys whatever |
| 81 | // metadata happened to be in the table source node's output schema). This more |
| 82 | // accurately simulates reading from a dataset. |
| 83 | acero::Declaration plan = acero::Declaration::Sequence( |
| 84 | {{"table_source", acero::TableSourceNodeOptions(table)}, |
| 85 | {"project", |
| 86 | acero::ProjectNodeOptions({compute::field_ref(0), compute::field_ref(1)})}, |
| 87 | {"write", write_options}}); |
| 88 | |
| 89 | ASSERT_OK(DeclarationToStatus(plan)); |
| 90 | |
| 91 | // Read the file back out and verify the nullability |
| 92 | ASSERT_OK_AND_ASSIGN(std::shared_ptr<io::RandomAccessFile> file, |
| 93 | mock_fs_->OpenInputFile("/my_dataset/0.arrow")); |
| 94 | ASSERT_OK_AND_ASSIGN(std::shared_ptr<ipc::RecordBatchFileReader> file_reader, |
| 95 | ipc::RecordBatchFileReader::Open(file)); |
| 96 | std::shared_ptr<Schema> file_schema = file_reader->schema(); |
| 97 | |
| 98 | ASSERT_TRUE(file_schema->field(0)->nullable()); |
| 99 | ASSERT_FALSE(file_schema->field(1)->nullable()); |
| 100 | |
| 101 | // Invalid custom schema |
| 102 | |
| 103 | // Incorrect # of fields |
| 104 | write_options.custom_schema = schema({}); |
| 105 | plan = acero::Declaration::Sequence( |
| 106 | {{"table_source", acero::TableSourceNodeOptions(table)}, |
| 107 | {"project", |
| 108 | acero::ProjectNodeOptions({compute::field_ref(0), compute::field_ref(1)})}, |
| 109 | {"write", write_options}}); |
| 110 | |
| 111 | ASSERT_THAT( |
| 112 | DeclarationToStatus(plan), |
| 113 | Raises(StatusCode::TypeError, |
| 114 | ::testing::HasSubstr("did not have the same number of fields as the data"))); |
| 115 | |
| 116 | // Incorrect types |
| 117 | write_options.custom_schema = |
| 118 | schema({field("nullable_i32", int32()), field("non_nullable_i32", int32())}); |
| 119 | plan = acero::Declaration::Sequence( |
nothing calls this directly
no test coverage detected