EnricoMi commented on code in PR #44470:
URL: https://github.com/apache/arrow/pull/44470#discussion_r1984739347
##########
cpp/src/arrow/dataset/write_node_test.cc:
##########
@@ -170,5 +176,83 @@ TEST_F(SimpleWriteNodeTest, CustomMetadata) {
ASSERT_TRUE(custom_metadata->Equals(*file_schema->metadata()));
}
+TEST_F(SimpleWriteNodeTest, SequenceOutput) {
+ // Test for ARROW-26818
+ auto format = std::make_shared<IpcFileFormat>();
+ constexpr int kRowsPerBatch = 16;
+ constexpr int kNumBatches = 32;
+ constexpr random::SeedType kSeed = 42;
+ constexpr int kJitterMod = 4;
+ acero::RegisterTestNodes();
+
+ // Create an input table
+ std::shared_ptr<Table> table =
+ gen::Gen({gen::Step()})->FailOnError()->Table(kRowsPerBatch,
kNumBatches);
+ auto dataset = std::make_shared<InMemoryDataset>(table);
+ auto scan_options = std::make_shared<ScanOptions>();
+ scan_options->use_threads = true;
+
+ for (bool preserve_order : {true, false}) {
+ auto scanner_builder = std::make_shared<ScannerBuilder>(dataset,
scan_options);
+ EXPECT_OK_AND_ASSIGN(auto scanner, scanner_builder->Finish());
+ auto exprs = scan_options->projection.call()->arguments;
+ auto names = checked_cast<const compute::MakeStructOptions*>(
+ scan_options->projection.call()->options.get())
+ ->field_names;
+
+ auto fs = std::make_shared<fs::internal::MockFileSystem>(fs::kNoTime);
+ dataset::WriteNodeOptions write_options(fs_write_options_);
+ write_options.write_options.file_write_options =
format->DefaultWriteOptions();
+ write_options.write_options.base_dir = "root";
+ write_options.write_options.partitioning =
+ std::make_shared<HivePartitioning>(schema({}));
+ write_options.write_options.basename_template = "{i}.feather";
+ write_options.write_options.filesystem = fs;
+ write_options.write_options.preserve_order = preserve_order;
+
+ // test plan of FileSystemDataset::Write with a jitter node that
guarantees exec
+ // batches are out of order
+ acero::Declaration plan = acero::Declaration::Sequence({
+ {"scan",
+ ScanNodeOptions{dataset, scanner->options(),
/*require_sequenced_output=*/false,
+ /*implicit_ordering=*/true}},
+ {"filter", acero::FilterNodeOptions{scanner->options()->filter}},
+ {"project", acero::ProjectNodeOptions{std::move(exprs),
std::move(names)}},
+ {"jitter", acero::JitterNodeOptions(kSeed, kJitterMod)},
+ {"write", write_options},
+ });
+
+ ASSERT_OK(DeclarationToStatus(plan));
+
+ // Read the file back out and verify the order
+ ASSERT_OK_AND_ASSIGN(auto dataset_factory, FileSystemDatasetFactory::Make(
+ fs, {"root/0.feather"},
format, {}));
+ ASSERT_OK_AND_ASSIGN(auto written_dataset,
dataset_factory->Finish(FinishOptions{}));
+ ASSERT_OK_AND_ASSIGN(scanner_builder, written_dataset->NewScan());
+ ASSERT_OK(scanner_builder->UseThreads(false));
+ ASSERT_OK_AND_ASSIGN(scanner, scanner_builder->Finish());
+ ASSERT_OK_AND_ASSIGN(auto actual, scanner->ToTable());
+ TableBatchReader reader(*actual);
+ std::shared_ptr<RecordBatch> batch;
+ ABORT_NOT_OK(reader.ReadNext(&batch));
Review Comment:
fixed
--
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-unsubscr...@arrow.apache.org
For queries about this service, please contact Infrastructure at:
us...@infra.apache.org
