GH-41149: [C++][Acero] Fix asof join race

cpp/src/arrow/acero/asof_join_node.cc

@@ -548,8 +548,10 @@ class InputState {
   // true when the queue is empty and, when memo may have future entries (the case of a
   // positive tolerance), when the memo is empty.
   // used when checking whether RHS is up to date with LHS.
-  bool CurrentEmpty() const {
-    return memo_.no_future_ ? Empty() : memo_.times_.empty() && Empty();
+  // NOTE: The emptiness must be decided by a single call to Empty() in caller, due to the
+  // potential race with Push(), see GH-41614.
+  bool CurrentEmpty(bool empty) const {
+    return memo_.no_future_ ? empty : (memo_.times_.empty() && empty);
   }
 
   // in case memo may not have future entries (the case of a non-positive tolerance),
@@ -650,13 +652,15 @@ class InputState {
   // timestamp, update latest_time and latest_ref_row to the value that immediately pass
   // the horizon. Update the memo-store with any entries or future entries so observed.
   // Returns true if updates were made, false if not.
-  Result<bool> AdvanceAndMemoize(OnType ts) {
+  // NOTE: The emptiness must be decided by a single call to Empty() in caller, due to the
+  // potential race with Push(), see GH-41614.
+  Result<bool> AdvanceAndMemoize(OnType ts, bool empty) {
     // Advance the right side row index until we reach the latest right row (for each key)
     // for the given left timestamp.
     DEBUG_SYNC(node_, "Advancing input ", index_, DEBUG_MANIP(std::endl));
 
     // Check if already updated for TS (or if there is no latest)
-    if (Empty()) {  // can't advance if empty and no future entries
+    if (empty) {  // can't advance if empty and no future entries
       return memo_.no_future_ ? false : memo_.RemoveEntriesWithLesserTime(ts);
     }
 
@@ -918,34 +922,46 @@ class CompositeTableBuilder {
 // guaranteeing this probability is below 1 in a billion. The fix is 128-bit hashing.
 // See ARROW-17653
 class AsofJoinNode : public ExecNode {
-  // Advances the RHS as far as possible to be up to date for the current LHS timestamp
-  Result<bool> UpdateRhs() {
+  // A simple wrapper for the result of a single call to UpdateRhs(), identifying:
+  // 1) If any RHS has advanced.
+  // 2) If all RHS are up to date with LHS.
+  struct RhsUpdateState {
+    bool any_advanced;
+    bool all_up_to_date_with_lhs;
+  };
+  // Advances the RHS as far as possible to be up to date for the current LHS timestamp,
+  // and checks if all RHS are up to date with LHS. The reason they have to be performed
+  // together is that they both depend on the emptiness of the RHS, which can be changed
+  // by Push() executing in another thread.
+  Result<RhsUpdateState> UpdateRhs() {
     auto& lhs = *state_.at(0);
     auto lhs_latest_time = lhs.GetLatestTime();
-    bool any_updated = false;
-    for (size_t i = 1; i < state_.size(); ++i) {
-      ARROW_ASSIGN_OR_RAISE(bool advanced, state_[i]->AdvanceAndMemoize(lhs_latest_time));
-      any_updated |= advanced;
-    }
-    return any_updated;
-  }
-
-  // Returns false if RHS not up to date for LHS
-  bool IsUpToDateWithLhsRow() const {
-    auto& lhs = *state_[0];
-    if (lhs.Empty()) return false;  // can't proceed if nothing on the LHS
-    OnType lhs_ts = lhs.GetLatestTime();
+    RhsUpdateState update_state{/*any_advanced=*/false, /*all_up_to_date_with_lhs=*/true};
     for (size_t i = 1; i < state_.size(); ++i) {
       auto& rhs = *state_[i];
-      if (!rhs.Finished()) {
+
+      // Obtain RHS emptiness once for subsequent AdvanceAndMemoize() and CurrentEmpty().
+      bool rhs_empty = rhs.Empty();
+      // Obtain RHS current time here because AdvanceAndMemoize() can change the
+      // emptiness.
+      OnType rhs_current_time = rhs_empty ? OnType{} : rhs.GetLatestTime();
+
+      ARROW_ASSIGN_OR_RAISE(bool advanced,
+                            rhs.AdvanceAndMemoize(lhs_latest_time, rhs_empty));
+      update_state.any_advanced |= advanced;
+
+      if (update_state.all_up_to_date_with_lhs && !rhs.Finished()) {
         // If RHS is finished, then we know it's up to date
-        if (rhs.CurrentEmpty())
-          return false;  // RHS isn't finished, but is empty --> not up to date
-        if (lhs_ts > rhs.GetCurrentTime())
-          return false;  // RHS isn't up to date (and not finished)
+        if (rhs.CurrentEmpty(rhs_empty)) {
+          // RHS isn't finished, but is empty --> not up to date
+          update_state.all_up_to_date_with_lhs = false;
+        } else if (lhs_latest_time > rhs_current_time) {
+          // RHS isn't up to date (and not finished)
+          update_state.all_up_to_date_with_lhs = false;
+        }
       }
     }
-    return true;
+    return update_state;
   }
 
   Result<std::shared_ptr<RecordBatch>> ProcessInner() {
@@ -963,20 +979,19 @@ class AsofJoinNode : public ExecNode {
       // If LHS is finished or empty then there's nothing we can do here
       if (lhs.Finished() || lhs.Empty()) break;
 
-      // Advance each of the RHS as far as possible to be up to date for the LHS timestamp
-      ARROW_ASSIGN_OR_RAISE(bool any_rhs_advanced, UpdateRhs());
+      ARROW_ASSIGN_OR_RAISE(auto rhs_update_state, UpdateRhs());
 
       // If we have received enough inputs to produce the next output batch
       // (decided by IsUpToDateWithLhsRow), we will perform the join and
       // materialize the output batch. The join is done by advancing through
       // the LHS and adding joined row to rows_ (done by Emplace). Finally,
       // input batches that are no longer needed are removed to free up memory.
-      if (IsUpToDateWithLhsRow()) {
+      if (rhs_update_state.all_up_to_date_with_lhs) {
         dst.Emplace(state_, tolerance_);
         ARROW_ASSIGN_OR_RAISE(bool advanced, lhs.Advance());
         if (!advanced) break;  // if we can't advance LHS, we're done for this batch
       } else {
-        if (!any_rhs_advanced) break;  // need to wait for new data
+        if (!rhs_update_state.any_advanced) break;  // need to wait for new data
       }
     }
 

cpp/src/arrow/acero/asof_join_node_test.cc

@@ -1678,5 +1678,59 @@ TEST(AsofJoinTest, BackpressureWithBatchesGen) {
                           /*slow_r0=*/false);
 }
 
+// Reproduction of GH-40675: A logical race between Process() and Push() that can be more
+// easily observed with single small batch.
+TEST(AsofJoinTest, RhsEmptinessRace) {
+  auto left_batch = ExecBatchFromJSON(
+      {int64(), utf8()}, R"([[1, "a"], [1, "b"], [5, "a"], [6, "b"], [7, "f"]])");
+  auto right_batch = ExecBatchFromJSON(
+      {int64(), utf8(), float64()}, R"([[2, "a", 1.0], [9, "b", 3.0], [15, "g", 5.0]])");
+
+  Declaration left{
+      "exec_batch_source",
+      ExecBatchSourceNodeOptions(schema({field("colA", int64()), field("col2", utf8())}),
+                                 {std::move(left_batch)})};
+  Declaration right{
+      "exec_batch_source",
+      ExecBatchSourceNodeOptions(schema({field("colB", int64()), field("col3", utf8()),
+                                         field("colC", float64())}),
+                                 {std::move(right_batch)})};
+  AsofJoinNodeOptions asof_join_opts({{{"colA"}, {{"col2"}}}, {{"colB"}, {{"col3"}}}}, 1);
+  Declaration asof_join{
+      "asofjoin", {std::move(left), std::move(right)}, std::move(asof_join_opts)};
+
+  ASSERT_OK_AND_ASSIGN(auto result, DeclarationToExecBatches(std::move(asof_join)));
+
+  auto exp_batch = ExecBatchFromJSON(
+      {int64(), utf8(), float64()},
+      R"([[1, "a", 1.0], [1, "b", null], [5, "a", null], [6, "b", null], [7, "f", null]])");
+  AssertExecBatchesEqualIgnoringOrder(result.schema, {exp_batch}, result.batches);
+}
+
+// Reproduction of GH-41149: Another case of the same root cause as GH-40675, but with
+// empty "by" columns.
+TEST(AsofJoinTest, RhsEmptinessRaceEmptyBy) {
+  auto left_batch = ExecBatchFromJSON({int64()}, R"([[1], [2], [3]])");
+  auto right_batch =
+      ExecBatchFromJSON({utf8(), int64()}, R"([["Z", 2], ["B", 3], ["A", 4]])");
+
+  Declaration left{"exec_batch_source",
+                   ExecBatchSourceNodeOptions(schema({field("on", int64())}),
+                                              {std::move(left_batch)})};
+  Declaration right{
+      "exec_batch_source",
+      ExecBatchSourceNodeOptions(schema({field("colVals", utf8()), field("on", int64())}),
+                                 {std::move(right_batch)})};
+  AsofJoinNodeOptions asof_join_opts({{{"on"}, {}}, {{"on"}, {}}}, 1);
+  Declaration asof_join{
+      "asofjoin", {std::move(left), std::move(right)}, std::move(asof_join_opts)};
+
+  ASSERT_OK_AND_ASSIGN(auto result, DeclarationToExecBatches(std::move(asof_join)));
+
+  auto exp_batch =
+      ExecBatchFromJSON({int64(), utf8()}, R"([[1, "Z"], [2, "Z"], [3, "B"]])");
+  AssertExecBatchesEqualIgnoringOrder(result.schema, {exp_batch}, result.batches);
+}
+
 }  // namespace acero
 }  // namespace arrow