zanmato1984 commented on code in PR #46455:
URL: https://github.com/apache/arrow/pull/46455#discussion_r3874125792
##########
cpp/src/arrow/dataset/file_test.cc:
##########
@@ -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<RecordBatch> 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<int32_t>(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<IpcFileFormat>();
+ auto fs = std::make_shared<fs::internal::MockFileSystem>(fs::kNoTime);
+ FileSystemDatasetWriteOptions write_options;
+ write_options.file_write_options = format->DefaultWriteOptions();
+ write_options.filesystem = fs;
+ write_options.partitioning = std::make_shared<HivePartitioning>(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<MockDataset>(schema({field("f0", int32())}));
+
+ auto delay_func = std::make_shared<compute::ScalarFunction>(
+ "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<std::optional<cp::ExecBatch>> 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);
Review Comment:
Non-blocking: Is this output-order check necessary for testing TeeNode?
The default SinkNode already sequences input with a meaningful ordering, so
the collected output should be ordered even if TeeNode forwards batches out of
order. As a result, this assertion does not directly exercise TeeNode's
`preserve_order` behavior. The checks against the datasets written by the
unordered and ordered TeeNodes below appear to provide the relevant coverage.
The sink is still needed to consume the plan, but the conversion to
`output_table` and its ordering assertion may be unnecessary.
##########
cpp/src/arrow/dataset/file_base.cc:
##########
@@ -559,13 +560,18 @@ Result<acero::ExecNode*> 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<acero::ExecNode*> inputs,
std::shared_ptr<Schema> 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) {
Review Comment:
Could we add the same input-ordering validation used by `SinkNode` and
`ConsumingSinkNode`?
With `preserve_order=true`, TeeNode creates a `SerialSequencingQueue` even
when `inputs_[0]->ordering()` is unordered. In that case, batches may retain
`kUnsequencedIndex`. `SerialSequencingQueue::InsertBatch` DCHECKs this in debug
builds; in release builds, such batches never match the queue's initial
`next_index_ == 0`, so they can remain queued and prevent the plan from
finishing.
TeeNode should reject this configuration in `Validate()`, with a regression
test for unordered input plus `preserve_order=true`.
##########
cpp/src/arrow/dataset/file_test.cc:
##########
@@ -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<RecordBatch> 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<int32_t>(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<IpcFileFormat>();
+ auto fs = std::make_shared<fs::internal::MockFileSystem>(fs::kNoTime);
+ FileSystemDatasetWriteOptions write_options;
+ write_options.file_write_options = format->DefaultWriteOptions();
+ write_options.filesystem = fs;
+ write_options.partitioning = std::make_shared<HivePartitioning>(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<MockDataset>(schema({field("f0", int32())}));
+
+ auto delay_func = std::make_shared<compute::ScalarFunction>(
Review Comment:
Non-blocking: Could this use the existing test `JitterNode` with a fixed
seed instead of relying on the scalar function's wall-clock delay?
`JitterNode` is designed to resequence physical batch delivery while
preserving batch indices, and is already used by the order-by tests for this
purpose. Reusing it here would make the test intent clearer and avoid depending
on thread timing and sleeps to produce out-of-order arrival.
--
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]
