zanmato1984 commented on code in PR #44470:
URL: https://github.com/apache/arrow/pull/44470#discussion_r1988460181
##########
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));
+ int32_t prev = -1;
+ auto out_of_order = false;
+ while (batch != nullptr) {
+ for (int row = 0; row < batch->num_rows(); ++row) {
+ auto scalar = batch->column(0)->GetScalar(row).ValueOrDie();
+ auto numeric_scalar =
+
std::static_pointer_cast<arrow::NumericScalar<arrow::Int32Type>>(scalar);
+ int32_t value = numeric_scalar->value;
+ if (value <= prev) {
+ out_of_order = true;
+ }
+ prev = value;
+ }
+ ABORT_NOT_OK(reader.ReadNext(&batch));
Review Comment:
This one is not updated yet.
--
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]
