From 37db2ccee4623afdb18cea87d9f106112e36fadd Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Rafa=C5=82=20Hibner?= Date: Thu, 15 May 2025 09:02:02 +0000 Subject: [PATCH 1/5] Preserve order when writting with TeeNode --- cpp/src/arrow/dataset/file_base.cc | 23 +++++++++++++++++++++-- 1 file changed, 21 insertions(+), 2 deletions(-) diff --git a/cpp/src/arrow/dataset/file_base.cc b/cpp/src/arrow/dataset/file_base.cc index ccc79dfa9bfc..0f1b8b52f5be 100644 --- a/cpp/src/arrow/dataset/file_base.cc +++ b/cpp/src/arrow/dataset/file_base.cc @@ -17,6 +17,7 @@ #include "arrow/dataset/file_base.h" +#include "arrow/acero/accumulation_queue.h" #include "arrow/acero/exec_plan.h" #include @@ -559,13 +560,18 @@ Result MakeWriteNode(acero::ExecPlan* plan, return node; } -class TeeNode : public acero::MapNode { +class TeeNode : public acero::MapNode, + public arrow::acero::util::SerialSequencingQueue::Processor { public: TeeNode(acero::ExecPlan* plan, std::vector inputs, std::shared_ptr output_schema, FileSystemDatasetWriteOptions write_options) : MapNode(plan, std::move(inputs), std::move(output_schema)), - write_options_(std::move(write_options)) {} + write_options_(std::move(write_options)) { + if (write_options.preserve_order) { + sequencer_ = acero::util::SerialSequencingQueue::Make(this); + } + } Status StartProducing() override { ARROW_ASSIGN_OR_RAISE( @@ -592,6 +598,18 @@ class TeeNode : public acero::MapNode { const char* kind_name() const override { return "TeeNode"; } + Status InputReceived(ExecNode* input, ExecBatch batch) override { + DCHECK_EQ(input, inputs_[0]); + if (sequencer_) { + return sequencer_->InsertBatch(std::move(batch)); + } + return Process(std::move(batch)); + } + + Status Process(ExecBatch batch) override { + return acero::MapNode::InputReceived(inputs_[0], batch); + } + void Finish() override { dataset_writer_->Finish(); } Result ProcessBatch(compute::ExecBatch batch) override { @@ -625,6 +643,7 @@ class TeeNode : public acero::MapNode { std::unique_ptr dataset_writer_; FileSystemDatasetWriteOptions write_options_; std::atomic backpressure_counter_ = 0; + std::unique_ptr sequencer_{nullptr}; }; } // namespace From 296726c03f6a5d0896597ffcadadbe9b0ad33728 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Rafa=C5=82=20Hibner?= Date: Wed, 26 Aug 2026 10:17:47 +0000 Subject: [PATCH 2/5] Add TeeNode ordering test --- cpp/src/arrow/dataset/file_test.cc | 118 +++++++++++++++++++++++++---- 1 file changed, 102 insertions(+), 16 deletions(-) diff --git a/cpp/src/arrow/dataset/file_test.cc b/cpp/src/arrow/dataset/file_test.cc index 2e2561203bea..54eec6066390 100644 --- a/cpp/src/arrow/dataset/file_test.cc +++ b/cpp/src/arrow/dataset/file_test.cc @@ -33,6 +33,7 @@ #include #include "arrow/acero/exec_plan.h" #include "arrow/acero/test_util_internal.h" +#include "arrow/acero/util.h" #include "arrow/array/array_primitive.h" #include "arrow/compute/test_util_internal.h" #include "arrow/dataset/api.h" @@ -444,6 +445,27 @@ class MockDataset : public Dataset { }; }; +Result HasOutOfOrderRows(const Table& table) { + TableBatchReader reader(table); + std::shared_ptr batch; + ARROW_RETURN_NOT_OK(reader.ReadNext(&batch)); + int32_t prev = 0; + bool has_prev = false; + while (batch != nullptr) { + const auto* values = batch->column(0)->data()->GetValues(1); + for (int row = 0; row < batch->num_rows(); ++row) { + int32_t value = values[row]; + if (has_prev && value <= prev) { + return true; + } + prev = value; + has_prev = true; + } + ARROW_RETURN_NOT_OK(reader.ReadNext(&batch)); + } + return false; +} + TEST_F(TestFileSystemDataset, MultiThreadedWritePersistsOrder) { // Test for GH-26818 // @@ -500,26 +522,90 @@ TEST_F(TestFileSystemDataset, MultiThreadedWritePersistsOrder) { 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 batch; - ASSERT_OK(reader.ReadNext(&batch)); - int32_t prev = -1; - auto out_of_order = false; - while (batch != nullptr) { - const auto* values = batch->column(0)->data()->GetValues(1); - for (int row = 0; row < batch->num_rows(); ++row) { - int32_t value = values[row]; - if (value <= prev) { - out_of_order = true; - } - prev = value; - } - ASSERT_OK(reader.ReadNext(&batch)); - } + ASSERT_OK_AND_ASSIGN(auto out_of_order, HasOutOfOrderRows(*actual)); ASSERT_EQ(!out_of_order, preserve_order); } } +TEST_F(TestFileSystemDataset, MultiThreadedTeeWritePersistsOrder) { + dataset::internal::Initialize(); + + auto format = std::make_shared(); + auto fs = std::make_shared(fs::kNoTime); + FileSystemDatasetWriteOptions write_options; + write_options.file_write_options = format->DefaultWriteOptions(); + write_options.filesystem = fs; + write_options.partitioning = std::make_shared(schema({})); + write_options.basename_template = "{i}.feather"; + + auto unordered_write_options = write_options; + unordered_write_options.base_dir = "unordered"; + unordered_write_options.preserve_order = false; + auto ordered_write_options = write_options; + ordered_write_options.base_dir = "ordered"; + ordered_write_options.preserve_order = true; + + auto dataset = std::make_shared(schema({field("f0", int32())})); + + auto delay_func = std::make_shared( + "tee_delay", compute::Arity(1), compute::FunctionDoc()); + compute::ScalarKernel delay_kernel; + delay_kernel.exec = delay; + delay_kernel.signature = compute::KernelSignature::Make({int32()}, boolean()); + ASSERT_OK(delay_func->AddKernel(delay_kernel)); + ASSERT_OK(compute::GetFunctionRegistry()->AddFunction(delay_func)); + + ASSERT_OK_AND_ASSIGN(auto scanner_builder, dataset->NewScan()); + ASSERT_OK(scanner_builder->UseThreads(true)); + ASSERT_OK( + scanner_builder->Filter(compute::call("tee_delay", {compute::field_ref("f0")}))); + ASSERT_OK_AND_ASSIGN(auto scanner, scanner_builder->Finish()); + + AsyncGenerator> sink_gen; + ASSERT_OK_AND_ASSIGN(auto plan, acero::ExecPlan::Make()); + // The first TeeNode records the delayed, out-of-order stream without changing it. + // The second TeeNode must use the batch indices to restore order. + ASSERT_OK( + acero::Declaration::Sequence( + { + {"scan", ScanNodeOptions{dataset, scanner->options(), + /*require_sequenced_output=*/true, + /*implicit_ordering=*/true}}, + {"filter", acero::FilterNodeOptions{scanner->options()->filter}}, + {"project", acero::ProjectNodeOptions{{compute::field_ref("f0")}, {"f0"}}}, + {"tee", WriteNodeOptions{unordered_write_options}, "unordered_tee"}, + {"tee", WriteNodeOptions{ordered_write_options}, "ordered_tee"}, + {"sink", acero::SinkNodeOptions{&sink_gen}}, + }) + .AddToPlan(plan.get())); + + ASSERT_FINISHES_OK_AND_ASSIGN(auto output_batches, + acero::StartAndCollect(plan.get(), sink_gen)); + ASSERT_OK_AND_ASSIGN(auto output_table, + acero::TableFromExecBatches(dataset->schema(), output_batches)); + ASSERT_OK_AND_ASSIGN(auto output_out_of_order, HasOutOfOrderRows(*output_table)); + ASSERT_FALSE(output_out_of_order); + + auto read_written_table = + [&](const std::string& path) -> Result> { + ARROW_ASSIGN_OR_RAISE(auto dataset_factory, + FileSystemDatasetFactory::Make(fs, {path}, format, {})); + ARROW_ASSIGN_OR_RAISE(auto written_dataset, dataset_factory->Finish(FinishOptions{})); + ARROW_ASSIGN_OR_RAISE(auto written_scanner_builder, written_dataset->NewScan()); + ARROW_RETURN_NOT_OK(written_scanner_builder->UseThreads(false)); + ARROW_ASSIGN_OR_RAISE(auto written_scanner, written_scanner_builder->Finish()); + return written_scanner->ToTable(); + }; + + ASSERT_OK_AND_ASSIGN(auto unordered_table, read_written_table("unordered/0.feather")); + ASSERT_OK_AND_ASSIGN(auto unordered_out_of_order, HasOutOfOrderRows(*unordered_table)); + ASSERT_TRUE(unordered_out_of_order); + + ASSERT_OK_AND_ASSIGN(auto ordered_table, read_written_table("ordered/0.feather")); + ASSERT_OK_AND_ASSIGN(auto ordered_out_of_order, HasOutOfOrderRows(*ordered_table)); + ASSERT_FALSE(ordered_out_of_order); +} + class FileSystemWriteTest : public testing::TestWithParam> { using PlanFactory = std::function( const FileSystemDatasetWriteOptions&, From 3fdbb7e9cd006f35e4e08e940b73abcc5c4580f1 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Rafa=C5=82=20Hibner?= Date: Fri, 28 Aug 2026 11:59:22 +0000 Subject: [PATCH 3/5] Add input ordering valiation --- cpp/src/arrow/dataset/file_base.cc | 10 +++++++++ cpp/src/arrow/dataset/file_test.cc | 33 ++++++++++++++++++++++++++++++ 2 files changed, 43 insertions(+) diff --git a/cpp/src/arrow/dataset/file_base.cc b/cpp/src/arrow/dataset/file_base.cc index 0f1b8b52f5be..8ad2254aa64b 100644 --- a/cpp/src/arrow/dataset/file_base.cc +++ b/cpp/src/arrow/dataset/file_base.cc @@ -598,6 +598,16 @@ class TeeNode : public acero::MapNode, const char* kind_name() const override { return "TeeNode"; } + Status Validate() const override { + ARROW_RETURN_NOT_OK(acero::MapNode::Validate()); + if (inputs_[0]->ordering().is_unordered() && sequencer_) { + return Status::Invalid("Tee node '", label(), + "' is configured to sequence output but there is no " + "meaningful ordering in the input"); + } + return Status::OK(); + } + Status InputReceived(ExecNode* input, ExecBatch batch) override { DCHECK_EQ(input, inputs_[0]); if (sequencer_) { diff --git a/cpp/src/arrow/dataset/file_test.cc b/cpp/src/arrow/dataset/file_test.cc index 54eec6066390..0c92365426aa 100644 --- a/cpp/src/arrow/dataset/file_test.cc +++ b/cpp/src/arrow/dataset/file_test.cc @@ -466,6 +466,39 @@ Result HasOutOfOrderRows(const Table& table) { return false; } +TEST_F(TestFileSystemDataset, RejectPreserveOrderWithUnorderedInput) { + dataset::internal::Initialize(); + + auto format = std::make_shared(); + FileSystemDatasetWriteOptions write_options; + write_options.file_write_options = format->DefaultWriteOptions(); + write_options.filesystem = std::make_shared(fs::kNoTime); + write_options.base_dir = "root"; + write_options.partitioning = std::make_shared(schema({})); + write_options.basename_template = "{i}.feather"; + write_options.preserve_order = true; + + auto source_data = acero::MakeBasicBatches(); + for (const char* factory_name : {"write", "tee"}) { + SCOPED_TRACE(factory_name); + ASSERT_OK_AND_ASSIGN(auto plan, acero::ExecPlan::Make()); + AsyncGenerator> sink_gen; + std::vector declarations = { + {"source", + acero::SourceNodeOptions{source_data.schema, source_data.gen(false, false)}}, + {factory_name, WriteNodeOptions{write_options}}, + }; + if (std::string(factory_name) == "tee") { + declarations.emplace_back("sink", acero::SinkNodeOptions{&sink_gen}); + } + ASSERT_OK( + acero::Declaration::Sequence(std::move(declarations)).AddToPlan(plan.get())); + ASSERT_THAT(plan->Validate(), + Raises(StatusCode::Invalid, + ::testing::HasSubstr("no meaningful ordering in the input"))); + } +} + TEST_F(TestFileSystemDataset, MultiThreadedWritePersistsOrder) { // Test for GH-26818 // From 18eee2078d4ce4e5e8ecc2cabe89c273a54e2b1f Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Rafa=C5=82=20Hibner?= Date: Fri, 28 Aug 2026 12:29:07 +0000 Subject: [PATCH 4/5] Use jitter node in tests --- cpp/src/arrow/dataset/file_test.cc | 188 +++++------------------------ 1 file changed, 31 insertions(+), 157 deletions(-) diff --git a/cpp/src/arrow/dataset/file_test.cc b/cpp/src/arrow/dataset/file_test.cc index 0c92365426aa..3da0eb7f190e 100644 --- a/cpp/src/arrow/dataset/file_test.cc +++ b/cpp/src/arrow/dataset/file_test.cc @@ -15,12 +15,9 @@ // specific language governing permissions and limitations // under the License. -#include -#include #include #include #include -#include #include #include @@ -32,8 +29,8 @@ #include #include #include "arrow/acero/exec_plan.h" +#include "arrow/acero/test_nodes.h" #include "arrow/acero/test_util_internal.h" -#include "arrow/acero/util.h" #include "arrow/array/array_primitive.h" #include "arrow/compute/test_util_internal.h" #include "arrow/dataset/api.h" @@ -45,6 +42,7 @@ #include "arrow/filesystem/test_util.h" #include "arrow/status.h" #include "arrow/testing/future_util.h" +#include "arrow/testing/generator.h" #include "arrow/testing/gtest_util.h" #include "arrow/util/io_util.h" @@ -362,88 +360,10 @@ TEST_F(TestFileSystemDataset, WriteProjected) { } } -// This kernel delays execution for some specific scalar values, -// which guarantees the writing phase sees out-of-order exec batches -Status delay(compute::KernelContext* ctx, const compute::ExecSpan& batch, - compute::ExecResult* out) { - const ArraySpan& input = batch[0].array; - const auto* input_values = input.GetValues(1); - uint8_t* output_values = out->array_span()->buffers[1].data; - - // Boolean data is stored in 1 bit per value - for (int64_t i = 0; i < input.length; ++i) { - if (input_values[i] % 16 == 0) { - std::this_thread::sleep_for(std::chrono::milliseconds(10)); - } - bit_util::SetBitTo(output_values, i, true); - } - - return Status::OK(); -} - -// A fragment with start=0 will defer ScanBatchesAsync returning a batch generator -// This guarantees a dataset of multiple fragments could produce out-of-order batches -class MockFragment : public Fragment { - public: - explicit MockFragment(uint32_t start, int64_t rows_per_batch, int num_batches, - const std::shared_ptr& schema) - : Fragment(compute::literal(true), schema), - start_(start), - rows_per_batch_(rows_per_batch), - num_batches_(num_batches) {} - - Result ScanBatchesAsync( - const std::shared_ptr& options) override { - // Fragment with start_=0 defers returning the generator - if (start_ == 0) { - std::this_thread::sleep_for(std::chrono::duration(0.1)); - } - - auto vec = gen::Gen({gen::Step(start_)}) - ->FailOnError() - ->RecordBatches(rows_per_batch_, num_batches_); - auto it = MakeVectorIterator(vec); - return MakeBackgroundGenerator(std::move(it), io::default_io_context().executor()); - } - - std::string type_name() const override { return "mock"; } - - protected: - Result> ReadPhysicalSchemaImpl() override { - return given_physical_schema_; - }; - - private: - uint32_t start_; - int64_t rows_per_batch_; - int num_batches_; -}; - -// This dataset consists of multiple fragments with incrementing values across the -// fragments -class MockDataset : public Dataset { - public: - explicit MockDataset(const std::shared_ptr& schema) : Dataset(schema) {} - - MockDataset(const std::shared_ptr& schema, - const compute::Expression& partition_expression) - : Dataset(schema, partition_expression) {} - - std::string type_name() const override { return "mock"; } - Result> ReplaceSchema( - std::shared_ptr schema) const override { - RETURN_NOT_OK(CheckProjectable(*schema_, *schema)); - return std::make_shared(std::move(schema)); - } - - protected: - Result GetFragmentsImpl(compute::Expression predicate) override { - FragmentVector fragments; - fragments.push_back(std::make_shared(0, 2, 1024, schema_)); - fragments.push_back(std::make_shared(2 * 1024, 2, 1024, schema_)); - return MakeVectorIterator(std::move(fragments)); - }; -}; +constexpr random::SeedType kJitterSeed = 42; +constexpr int kMaxJitterModifier = 4; +constexpr int64_t kOrderingRowsPerBatch = 1; +constexpr int kOrderingNumBatches = 256; Result HasOutOfOrderRows(const Table& table) { TableBatchReader reader(table); @@ -502,17 +422,12 @@ TEST_F(TestFileSystemDataset, RejectPreserveOrderWithUnorderedInput) { TEST_F(TestFileSystemDataset, MultiThreadedWritePersistsOrder) { // Test for GH-26818 // - // This test uses std::this_thread::sleep_for to increase chances for batches - // to get written out-of-order in multi-threaded environment. - // With preserve_order = false, the existence of out-of-order is asserted to - // verify that the test setup reliably writes out-of-order sequences, and - // that write_options.preserve_order = preserve_order can recreate order. - // - // Estimates for out_of_order == false and preserve_order == false to occur - // are 10^-62 https://github.com/apache/arrow/pull/44470#discussion_r2079049038 - // - // If this test starts to reliably fail with preserve_order == false, the test setup - // has to be revised to again reliably produce out-of-order sequences. + // JitterNode changes physical batch delivery order while preserving the meaningful + // batch indices assigned by TableSourceNode. The unordered write verifies the test + // setup, and the ordered write verifies that WriteNode restores the indexed order. + dataset::internal::Initialize(); + acero::RegisterTestNodes(); + auto format = std::make_shared(); FileSystemDatasetWriteOptions write_options; write_options.file_write_options = format->DefaultWriteOptions(); @@ -520,40 +435,27 @@ TEST_F(TestFileSystemDataset, MultiThreadedWritePersistsOrder) { write_options.partitioning = std::make_shared(schema({})); write_options.basename_template = "{i}.feather"; - // The Mock dataset delays emitting the first fragment, which test sequenced output of - // scan node - auto dataset = std::make_shared(schema({field("f0", int32())})); - - // The delay scalar function delays some batches of all fragments, which tests implicit - // ordering - auto delay_func = std::make_shared("delay", compute::Arity(1), - compute::FunctionDoc()); - compute::ScalarKernel delay_kernel; - delay_kernel.exec = delay; - delay_kernel.signature = compute::KernelSignature::Make({int32()}, boolean()); - ASSERT_OK(delay_func->AddKernel(delay_kernel)); - ASSERT_OK(compute::GetFunctionRegistry()->AddFunction(delay_func)); + auto input = gen::Gen({gen::Step()}) + ->FailOnError() + ->Table(kOrderingRowsPerBatch, kOrderingNumBatches); for (bool preserve_order : {true, false}) { - ASSERT_OK_AND_ASSIGN(auto scanner_builder, dataset->NewScan()); - ASSERT_OK(scanner_builder->UseThreads(true)); - ASSERT_OK( - scanner_builder->Filter(compute::call("delay", {compute::field_ref("f0")}))); - ASSERT_OK_AND_ASSIGN(auto scanner, scanner_builder->Finish()); - auto fs = std::make_shared(fs::kNoTime); write_options.filesystem = fs; write_options.preserve_order = preserve_order; - ASSERT_OK(FileSystemDataset::Write(write_options, scanner)); + ASSERT_OK(acero::DeclarationToStatus(acero::Declaration::Sequence( + {{"table_source", acero::TableSourceNodeOptions{input}}, + {"jitter", acero::JitterNodeOptions{kJitterSeed, kMaxJitterModifier}}, + {"write", WriteNodeOptions{write_options}}}))); // 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_AND_ASSIGN(auto scanner_builder, written_dataset->NewScan()); ASSERT_OK(scanner_builder->UseThreads(false)); - ASSERT_OK_AND_ASSIGN(scanner, scanner_builder->Finish()); + ASSERT_OK_AND_ASSIGN(auto scanner, scanner_builder->Finish()); ASSERT_OK_AND_ASSIGN(auto actual, scanner->ToTable()); ASSERT_OK_AND_ASSIGN(auto out_of_order, HasOutOfOrderRows(*actual)); ASSERT_EQ(!out_of_order, preserve_order); @@ -562,6 +464,7 @@ TEST_F(TestFileSystemDataset, MultiThreadedWritePersistsOrder) { TEST_F(TestFileSystemDataset, MultiThreadedTeeWritePersistsOrder) { dataset::internal::Initialize(); + acero::RegisterTestNodes(); auto format = std::make_shared(); auto fs = std::make_shared(fs::kNoTime); @@ -578,46 +481,17 @@ TEST_F(TestFileSystemDataset, MultiThreadedTeeWritePersistsOrder) { ordered_write_options.base_dir = "ordered"; ordered_write_options.preserve_order = true; - auto dataset = std::make_shared(schema({field("f0", int32())})); - - auto delay_func = std::make_shared( - "tee_delay", compute::Arity(1), compute::FunctionDoc()); - compute::ScalarKernel delay_kernel; - delay_kernel.exec = delay; - delay_kernel.signature = compute::KernelSignature::Make({int32()}, boolean()); - ASSERT_OK(delay_func->AddKernel(delay_kernel)); - ASSERT_OK(compute::GetFunctionRegistry()->AddFunction(delay_func)); - - ASSERT_OK_AND_ASSIGN(auto scanner_builder, dataset->NewScan()); - ASSERT_OK(scanner_builder->UseThreads(true)); - ASSERT_OK( - scanner_builder->Filter(compute::call("tee_delay", {compute::field_ref("f0")}))); - ASSERT_OK_AND_ASSIGN(auto scanner, scanner_builder->Finish()); + auto input = gen::Gen({gen::Step()}) + ->FailOnError() + ->Table(kOrderingRowsPerBatch, kOrderingNumBatches); - AsyncGenerator> sink_gen; - ASSERT_OK_AND_ASSIGN(auto plan, acero::ExecPlan::Make()); - // The first TeeNode records the delayed, out-of-order stream without changing it. + // The first TeeNode records the jittered, out-of-order stream without changing it. // The second TeeNode must use the batch indices to restore order. - ASSERT_OK( - acero::Declaration::Sequence( - { - {"scan", ScanNodeOptions{dataset, scanner->options(), - /*require_sequenced_output=*/true, - /*implicit_ordering=*/true}}, - {"filter", acero::FilterNodeOptions{scanner->options()->filter}}, - {"project", acero::ProjectNodeOptions{{compute::field_ref("f0")}, {"f0"}}}, - {"tee", WriteNodeOptions{unordered_write_options}, "unordered_tee"}, - {"tee", WriteNodeOptions{ordered_write_options}, "ordered_tee"}, - {"sink", acero::SinkNodeOptions{&sink_gen}}, - }) - .AddToPlan(plan.get())); - - ASSERT_FINISHES_OK_AND_ASSIGN(auto output_batches, - acero::StartAndCollect(plan.get(), sink_gen)); - ASSERT_OK_AND_ASSIGN(auto output_table, - acero::TableFromExecBatches(dataset->schema(), output_batches)); - ASSERT_OK_AND_ASSIGN(auto output_out_of_order, HasOutOfOrderRows(*output_table)); - ASSERT_FALSE(output_out_of_order); + ASSERT_OK(acero::DeclarationToStatus(acero::Declaration::Sequence( + {{"table_source", acero::TableSourceNodeOptions{input}}, + {"jitter", acero::JitterNodeOptions{kJitterSeed, kMaxJitterModifier}}, + {"tee", WriteNodeOptions{unordered_write_options}, "unordered_tee"}, + {"tee", WriteNodeOptions{ordered_write_options}, "ordered_tee"}}))); auto read_written_table = [&](const std::string& path) -> Result> { From e3fadedd6794b8a240c8c5b61e0dc216366103fd Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Rafa=C5=82=20Hibner?= Date: Fri, 28 Aug 2026 15:54:00 +0000 Subject: [PATCH 5/5] restore original MultiThreadedWritePersistsOrder --- cpp/src/arrow/dataset/file_test.cc | 133 ++++++++++++++++++++++++++--- 1 file changed, 119 insertions(+), 14 deletions(-) diff --git a/cpp/src/arrow/dataset/file_test.cc b/cpp/src/arrow/dataset/file_test.cc index 3da0eb7f190e..bd87a58eb41d 100644 --- a/cpp/src/arrow/dataset/file_test.cc +++ b/cpp/src/arrow/dataset/file_test.cc @@ -15,9 +15,12 @@ // specific language governing permissions and limitations // under the License. +#include +#include #include #include #include +#include #include #include @@ -360,6 +363,89 @@ TEST_F(TestFileSystemDataset, WriteProjected) { } } +// This kernel delays execution for some specific scalar values, +// which guarantees the writing phase sees out-of-order exec batches +Status delay(compute::KernelContext* ctx, const compute::ExecSpan& batch, + compute::ExecResult* out) { + const ArraySpan& input = batch[0].array; + const auto* input_values = input.GetValues(1); + uint8_t* output_values = out->array_span()->buffers[1].data; + + // Boolean data is stored in 1 bit per value + for (int64_t i = 0; i < input.length; ++i) { + if (input_values[i] % 16 == 0) { + std::this_thread::sleep_for(std::chrono::milliseconds(10)); + } + bit_util::SetBitTo(output_values, i, true); + } + + return Status::OK(); +} + +// A fragment with start=0 will defer ScanBatchesAsync returning a batch generator +// This guarantees a dataset of multiple fragments could produce out-of-order batches +class MockFragment : public Fragment { + public: + explicit MockFragment(uint32_t start, int64_t rows_per_batch, int num_batches, + const std::shared_ptr& schema) + : Fragment(compute::literal(true), schema), + start_(start), + rows_per_batch_(rows_per_batch), + num_batches_(num_batches) {} + + Result ScanBatchesAsync( + const std::shared_ptr& options) override { + // Fragment with start_=0 defers returning the generator + if (start_ == 0) { + std::this_thread::sleep_for(std::chrono::duration(0.1)); + } + + auto vec = gen::Gen({gen::Step(start_)}) + ->FailOnError() + ->RecordBatches(rows_per_batch_, num_batches_); + auto it = MakeVectorIterator(vec); + return MakeBackgroundGenerator(std::move(it), io::default_io_context().executor()); + } + + std::string type_name() const override { return "mock"; } + + protected: + Result> ReadPhysicalSchemaImpl() override { + return given_physical_schema_; + }; + + private: + uint32_t start_; + int64_t rows_per_batch_; + int num_batches_; +}; + +// This dataset consists of multiple fragments with incrementing values across the +// fragments +class MockDataset : public Dataset { + public: + explicit MockDataset(const std::shared_ptr& schema) : Dataset(schema) {} + + MockDataset(const std::shared_ptr& schema, + const compute::Expression& partition_expression) + : Dataset(schema, partition_expression) {} + + std::string type_name() const override { return "mock"; } + Result> ReplaceSchema( + std::shared_ptr schema) const override { + RETURN_NOT_OK(CheckProjectable(*schema_, *schema)); + return std::make_shared(std::move(schema)); + } + + protected: + Result GetFragmentsImpl(compute::Expression predicate) override { + FragmentVector fragments; + fragments.push_back(std::make_shared(0, 2, 1024, schema_)); + fragments.push_back(std::make_shared(2 * 1024, 2, 1024, schema_)); + return MakeVectorIterator(std::move(fragments)); + }; +}; + constexpr random::SeedType kJitterSeed = 42; constexpr int kMaxJitterModifier = 4; constexpr int64_t kOrderingRowsPerBatch = 1; @@ -422,11 +508,17 @@ TEST_F(TestFileSystemDataset, RejectPreserveOrderWithUnorderedInput) { TEST_F(TestFileSystemDataset, MultiThreadedWritePersistsOrder) { // Test for GH-26818 // - // JitterNode changes physical batch delivery order while preserving the meaningful - // batch indices assigned by TableSourceNode. The unordered write verifies the test - // setup, and the ordered write verifies that WriteNode restores the indexed order. - dataset::internal::Initialize(); - acero::RegisterTestNodes(); + // This test uses std::this_thread::sleep_for to increase chances for batches + // to get written out-of-order in multi-threaded environment. + // With preserve_order = false, the existence of out-of-order is asserted to + // verify that the test setup reliably writes out-of-order sequences, and + // that write_options.preserve_order = preserve_order can recreate order. + // + // Estimates for out_of_order == false and preserve_order == false to occur + // are 10^-62 https://github.com/apache/arrow/pull/44470#discussion_r2079049038 + // + // If this test starts to reliably fail with preserve_order == false, the test setup + // has to be revised to again reliably produce out-of-order sequences. auto format = std::make_shared(); FileSystemDatasetWriteOptions write_options; @@ -435,27 +527,40 @@ TEST_F(TestFileSystemDataset, MultiThreadedWritePersistsOrder) { write_options.partitioning = std::make_shared(schema({})); write_options.basename_template = "{i}.feather"; - auto input = gen::Gen({gen::Step()}) - ->FailOnError() - ->Table(kOrderingRowsPerBatch, kOrderingNumBatches); + // The Mock dataset delays emitting the first fragment, which test sequenced output of + // scan node + auto dataset = std::make_shared(schema({field("f0", int32())})); + + // The delay scalar function delays some batches of all fragments, which tests implicit + // ordering + auto delay_func = std::make_shared("delay", compute::Arity(1), + compute::FunctionDoc()); + compute::ScalarKernel delay_kernel; + delay_kernel.exec = delay; + delay_kernel.signature = compute::KernelSignature::Make({int32()}, boolean()); + ASSERT_OK(delay_func->AddKernel(delay_kernel)); + ASSERT_OK(compute::GetFunctionRegistry()->AddFunction(delay_func)); for (bool preserve_order : {true, false}) { + ASSERT_OK_AND_ASSIGN(auto scanner_builder, dataset->NewScan()); + ASSERT_OK(scanner_builder->UseThreads(true)); + ASSERT_OK( + scanner_builder->Filter(compute::call("delay", {compute::field_ref("f0")}))); + ASSERT_OK_AND_ASSIGN(auto scanner, scanner_builder->Finish()); + auto fs = std::make_shared(fs::kNoTime); write_options.filesystem = fs; write_options.preserve_order = preserve_order; - ASSERT_OK(acero::DeclarationToStatus(acero::Declaration::Sequence( - {{"table_source", acero::TableSourceNodeOptions{input}}, - {"jitter", acero::JitterNodeOptions{kJitterSeed, kMaxJitterModifier}}, - {"write", WriteNodeOptions{write_options}}}))); + ASSERT_OK(FileSystemDataset::Write(write_options, scanner)); // 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(auto scanner_builder, written_dataset->NewScan()); + ASSERT_OK_AND_ASSIGN(scanner_builder, written_dataset->NewScan()); ASSERT_OK(scanner_builder->UseThreads(false)); - ASSERT_OK_AND_ASSIGN(auto scanner, scanner_builder->Finish()); + ASSERT_OK_AND_ASSIGN(scanner, scanner_builder->Finish()); ASSERT_OK_AND_ASSIGN(auto actual, scanner->ToTable()); ASSERT_OK_AND_ASSIGN(auto out_of_order, HasOutOfOrderRows(*actual)); ASSERT_EQ(!out_of_order, preserve_order);