Skip to content

Commit c22181f

Browse files
author
Rafał Hibner
committed
Add error reporting tests
1 parent 596a60f commit c22181f

3 files changed

Lines changed: 191 additions & 11 deletions

File tree

‎cpp/src/arrow/acero/plan_test.cc‎

Lines changed: 167 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1788,5 +1788,172 @@ TEST(ExecPlanExecution, UnalignedInput) {
17881788
ASSERT_LT(initial_bytes_allocated, default_memory_pool()->total_bytes_allocated());
17891789
}
17901790

1791+
struct ExecPlanErrorReporting : public testing::TestWithParam<DummyNodeStatusReporter> {};
1792+
1793+
#pragma GCC push_options
1794+
#pragma GCC optimize("O0")
1795+
TEST_P(ExecPlanErrorReporting, SourceSink) {
1796+
ASSERT_OK_AND_ASSIGN(auto plan, ExecPlan::Make());
1797+
auto source = MakeDummyNode(plan.get(), "source", /*inputs=*/{}, /*is_sink=*/false,
1798+
/*start_producing =*/{},
1799+
/*stop_producing =*/{},
1800+
/*status_reporter =*/GetParam());
1801+
auto sink = MakeDummyNode(plan.get(), "sink", /*inputs=*/{source}, /*is_sink=*/true,
1802+
/*start_producing =*/{},
1803+
/*stop_producing =*/{},
1804+
/*status_reporter =*/GetParam());
1805+
1806+
ASSERT_OK(plan->Validate());
1807+
EXPECT_THAT(plan->nodes(), ElementsAre(source, sink));
1808+
1809+
bool should_finish = GetParam().start_producing.ok();
1810+
plan->StartProducing();
1811+
SleepABit();
1812+
if (should_finish)
1813+
ASSERT_FINISHES_OK(plan->finished());
1814+
else
1815+
ASSERT_FINISHES_AND_RAISES(Invalid, plan->finished());
1816+
}
1817+
1818+
TEST_P(ExecPlanErrorReporting, InputReceived) {
1819+
ASSERT_OK_AND_ASSIGN(auto plan, ExecPlan::Make());
1820+
auto basic_data = MakeBasicBatches();
1821+
1822+
ASSERT_OK_AND_ASSIGN(
1823+
auto source,
1824+
Declaration("source",
1825+
SourceNodeOptions{basic_data.schema, basic_data.gen(/*parallel=*/false,
1826+
/*slow=*/false)})
1827+
.AddToPlan(plan.get()));
1828+
1829+
MakeDummyNode(plan.get(), "sink", /*inputs=*/{source}, /*is_sink=*/true,
1830+
/*start_producing =*/{},
1831+
/*stop_producing =*/{},
1832+
/*status_reporter =*/GetParam());
1833+
1834+
bool should_finish = GetParam().start_producing.ok() &&
1835+
GetParam().input_received.ok() && GetParam().input_finished.ok();
1836+
plan->StartProducing();
1837+
SleepABit();
1838+
if (should_finish)
1839+
ASSERT_FINISHES_OK(plan->finished());
1840+
else
1841+
ASSERT_FINISHES_AND_RAISES(Invalid, plan->finished());
1842+
}
1843+
1844+
TEST_P(ExecPlanErrorReporting, Finish) {
1845+
std::shared_ptr<Schema> schema_ = schema({field("data", uint32())});
1846+
std::optional<ExecBatch> batch =
1847+
ExecBatchFromJSON({int32(), boolean()},
1848+
"[[4, false], [5, null], [6, false], [7, false], [null, true]]");
1849+
PushGenerator<std::optional<ExecBatch>> batch_producer;
1850+
ASSERT_OK_AND_ASSIGN(auto plan, ExecPlan::Make());
1851+
auto basic_data = MakeBasicBatches();
1852+
1853+
ASSERT_OK_AND_ASSIGN(auto source,
1854+
Declaration("source", SourceNodeOptions(schema_, batch_producer))
1855+
.AddToPlan(plan.get()));
1856+
1857+
MakeDummyNode(plan.get(), "sink", /*inputs=*/{source}, /*is_sink=*/true,
1858+
/*start_producing =*/{},
1859+
/*stop_producing =*/{},
1860+
/*status_reporter =*/GetParam());
1861+
1862+
bool should_start = GetParam().start_producing.ok();
1863+
plan->StartProducing();
1864+
SleepABit();
1865+
if (should_start) {
1866+
ASSERT_FALSE(plan->finished().is_finished());
1867+
} else {
1868+
ASSERT_FINISHES_AND_RAISES(Invalid, plan->finished());
1869+
}
1870+
batch_producer.producer().Push(batch);
1871+
SleepABit();
1872+
1873+
bool should_receive = should_start && GetParam().input_received.ok();
1874+
if (should_receive) {
1875+
ASSERT_FALSE(plan->finished().is_finished());
1876+
} else {
1877+
ASSERT_FINISHES_AND_RAISES(Invalid, plan->finished());
1878+
}
1879+
1880+
batch_producer.producer().Push(std::nullopt);
1881+
SleepABit();
1882+
bool should_finish = should_receive && GetParam().input_finished.ok();
1883+
if (should_finish) {
1884+
ASSERT_FINISHES_OK(plan->finished());
1885+
} else {
1886+
ASSERT_FINISHES_AND_RAISES(Invalid, plan->finished());
1887+
return;
1888+
}
1889+
}
1890+
1891+
TEST_P(ExecPlanErrorReporting, StopProducing) {
1892+
std::shared_ptr<Schema> schema_ = schema({field("data", uint32())});
1893+
std::optional<ExecBatch> batch =
1894+
ExecBatchFromJSON({int32(), boolean()},
1895+
"[[4, false], [5, null], [6, false], [7, false], [null, true]]");
1896+
PushGenerator<std::optional<ExecBatch>> batch_producer;
1897+
ASSERT_OK_AND_ASSIGN(auto plan, ExecPlan::Make());
1898+
auto basic_data = MakeBasicBatches();
1899+
1900+
ASSERT_OK_AND_ASSIGN(auto source,
1901+
Declaration("source", SourceNodeOptions(schema_, batch_producer))
1902+
.AddToPlan(plan.get()));
1903+
1904+
MakeDummyNode(plan.get(), "sink", /*inputs=*/{source}, /*is_sink=*/true,
1905+
/*start_producing =*/{},
1906+
/*stop_producing =*/{},
1907+
/*status_reporter =*/GetParam());
1908+
1909+
bool should_start = GetParam().start_producing.ok();
1910+
plan->StartProducing();
1911+
SleepABit();
1912+
if (should_start) {
1913+
ASSERT_FALSE(plan->finished().is_finished());
1914+
} else {
1915+
ASSERT_FINISHES_AND_RAISES(Invalid, plan->finished());
1916+
}
1917+
batch_producer.producer().Push(batch);
1918+
SleepABit();
1919+
1920+
bool should_receive = should_start && GetParam().input_received.ok();
1921+
if (should_receive) {
1922+
ASSERT_FALSE(plan->finished().is_finished());
1923+
} else {
1924+
ASSERT_FINISHES_AND_RAISES(Invalid, plan->finished());
1925+
}
1926+
// this should not be needed
1927+
batch_producer.producer().Push(std::nullopt);
1928+
SleepABit();
1929+
plan->StopProducing();
1930+
SleepABit();
1931+
bool should_stop =
1932+
should_receive && GetParam().stop_producing.ok() && GetParam().input_finished.ok();
1933+
if (should_stop) {
1934+
ASSERT_FINISHES_AND_RAISES(Cancelled, plan->finished());
1935+
} else {
1936+
ASSERT_FINISHES_AND_RAISES(Invalid, plan->finished());
1937+
return;
1938+
}
1939+
}
1940+
1941+
INSTANTIATE_TEST_SUITE_P(
1942+
ExecPlan, ExecPlanErrorReporting,
1943+
testing::Values(DummyNodeStatusReporter{Status::OK(), Status::OK(), Status::OK(),
1944+
Status::OK()},
1945+
DummyNodeStatusReporter{Status::Invalid("1"), Status::OK(),
1946+
Status::OK(), Status::OK()},
1947+
DummyNodeStatusReporter{Status::OK(), Status::Invalid("1"),
1948+
Status::OK(), Status::OK()},
1949+
DummyNodeStatusReporter{Status::OK(), Status::OK(),
1950+
Status::Invalid("1"), Status::OK()},
1951+
DummyNodeStatusReporter{Status::OK(), Status::OK(), Status::OK(),
1952+
Status::Invalid("1")},
1953+
DummyNodeStatusReporter{Status::Invalid("1"), Status::Invalid("1"),
1954+
Status::Invalid("1"), Status::Invalid("1")}));
1955+
1956+
#pragma GCC pop_options
1957+
17911958
} // namespace acero
17921959
} // namespace arrow

‎cpp/src/arrow/acero/test_util_internal.cc‎

Lines changed: 16 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -75,10 +75,12 @@ namespace {
7575

7676
struct DummyNode : ExecNode {
7777
DummyNode(ExecPlan* plan, NodeVector inputs, bool is_sink,
78-
StartProducingFunc start_producing, StopProducingFunc stop_producing)
78+
StartProducingFunc start_producing, StopProducingFunc stop_producing,
79+
DummyNodeStatusReporter status_reporter)
7980
: ExecNode(plan, std::move(inputs), {}, (is_sink) ? nullptr : dummy_schema()),
8081
start_producing_(std::move(start_producing)),
81-
stop_producing_(std::move(stop_producing)) {
82+
stop_producing_(std::move(stop_producing)),
83+
status_reporter_(status_reporter) {
8284
input_labels_.resize(inputs_.size());
8385
for (size_t i = 0; i < input_labels_.size(); ++i) {
8486
input_labels_[i] = std::to_string(i);
@@ -87,18 +89,20 @@ struct DummyNode : ExecNode {
8789

8890
const char* kind_name() const override { return "Dummy"; }
8991

90-
Status InputReceived(ExecNode* input, ExecBatch batch) override { return Status::OK(); }
92+
Status InputReceived(ExecNode* input, ExecBatch batch) override {
93+
return status_reporter_.input_received;
94+
}
9195

9296
Status InputFinished(ExecNode* input, int total_batches) override {
93-
return Status::OK();
97+
return status_reporter_.input_finished;
9498
}
9599

96100
Status StartProducing() override {
97101
if (start_producing_) {
98102
RETURN_NOT_OK(start_producing_(this));
99103
}
100104
started_ = true;
101-
return Status::OK();
105+
return status_reporter_.start_producing;
102106
}
103107

104108
void PauseProducing(ExecNode* output, int32_t counter) override {
@@ -115,7 +119,7 @@ struct DummyNode : ExecNode {
115119
if (stop_producing_) {
116120
stop_producing_(this);
117121
}
118-
return Status::OK();
122+
return status_reporter_.stop_producing;
119123
}
120124

121125
private:
@@ -127,6 +131,7 @@ struct DummyNode : ExecNode {
127131

128132
StartProducingFunc start_producing_;
129133
StopProducingFunc stop_producing_;
134+
DummyNodeStatusReporter status_reporter_;
130135
std::unordered_set<ExecNode*> requested_stop_;
131136
bool started_ = false;
132137
};
@@ -135,10 +140,11 @@ struct DummyNode : ExecNode {
135140

136141
ExecNode* MakeDummyNode(ExecPlan* plan, std::string label, std::vector<ExecNode*> inputs,
137142
bool is_sink, StartProducingFunc start_producing,
138-
StopProducingFunc stop_producing) {
139-
auto node =
140-
plan->EmplaceNode<DummyNode>(plan, std::move(inputs), is_sink,
141-
std::move(start_producing), std::move(stop_producing));
143+
StopProducingFunc stop_producing,
144+
DummyNodeStatusReporter status_reporter) {
145+
auto node = plan->EmplaceNode<DummyNode>(
146+
plan, std::move(inputs), is_sink, std::move(start_producing),
147+
std::move(stop_producing), std::move(status_reporter));
142148
if (!label.empty()) {
143149
node->SetLabel(std::move(label));
144150
}

‎cpp/src/arrow/acero/test_util_internal.h‎

Lines changed: 8 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -45,9 +45,16 @@ using StartProducingFunc = std::function<Status(ExecNode*)>;
4545
using StopProducingFunc = std::function<void(ExecNode*)>;
4646

4747
// Make a dummy node that has no execution behaviour
48+
struct DummyNodeStatusReporter {
49+
Status input_received;
50+
Status input_finished;
51+
Status start_producing;
52+
Status stop_producing;
53+
};
4854
ExecNode* MakeDummyNode(ExecPlan* plan, std::string label, std::vector<ExecNode*> inputs,
4955
bool is_sink = false, StartProducingFunc = {},
50-
StopProducingFunc = {});
56+
StopProducingFunc = {},
57+
DummyNodeStatusReporter status_reporter = {});
5158

5259
struct BatchesWithSchema {
5360
std::vector<ExecBatch> batches;

0 commit comments

Comments
 (0)