diff --git a/cpp/src/arrow/acero/asof_join_benchmark.cc b/cpp/src/arrow/acero/asof_join_benchmark.cc index ed2ac2258eb6..076a44fa4732 100644 --- a/cpp/src/arrow/acero/asof_join_benchmark.cc +++ b/cpp/src/arrow/acero/asof_join_benchmark.cc @@ -15,14 +15,21 @@ // specific language governing permissions and limitations // under the License. +#include +#include #include +#include +#include #include "benchmark/benchmark.h" #include "arrow/acero/options.h" #include "arrow/acero/test_util_internal.h" +#include "arrow/array/array_primitive.h" +#include "arrow/array/builder_binary.h" #include "arrow/table.h" #include "arrow/testing/future_util.h" +#include "arrow/util/byte_size.h" namespace arrow { namespace acero { @@ -33,6 +40,7 @@ const int kDefaultStart = 0; const int kDefaultEnd = 32000; const int kDefaultMinColumnVal = -10000; const int kDefaultMaxColumnVal = 10000; +const int64_t kLongStringKeyBytes = 128; struct TableStats { std::shared_ptr table; @@ -40,13 +48,50 @@ struct TableStats { size_t bytes; }; -static Result MakeTable(const TableGenerationProperties& properties) { +static Result> WithFixedSizeStringKeys( + std::shared_ptr
table, int num_ids, int64_t key_bytes) { + if (key_bytes == 0) { + return table; + } + + std::vector keys; + keys.reserve(num_ids); + for (int id = 0; id < num_ids; ++id) { + std::string suffix = std::to_string(id); + if (static_cast(suffix.size()) > key_bytes) { + return Status::Invalid("Key size is too small for id ", id); + } + keys.emplace_back(static_cast(key_bytes - suffix.size()), 'k'); + keys.back() += suffix; + } + + StringBuilder builder; + ARROW_RETURN_NOT_OK(builder.Resize(table->num_rows())); + ARROW_RETURN_NOT_OK(builder.ReserveData(table->num_rows() * key_bytes)); + for (const auto& chunk : table->GetColumnByName(kKeyCol)->chunks()) { + const auto& ids = static_cast(*chunk); + for (int64_t row = 0; row < ids.length(); ++row) { + builder.UnsafeAppend(keys[ids.Value(row)]); + } + } + + std::shared_ptr string_keys; + ARROW_RETURN_NOT_OK(builder.Finish(&string_keys)); + int key_index = table->schema()->GetFieldIndex(kKeyCol); + return table->SetColumn(key_index, field(kKeyCol, utf8()), + std::make_shared(std::move(string_keys))); +} + +static Result MakeTable(const TableGenerationProperties& properties, + int64_t string_key_bytes) { ARROW_ASSIGN_OR_RAISE(std::shared_ptr
table, MakeRandomTimeSeriesTable(properties)); - size_t row_size = sizeof(double) * (table.get()->schema()->num_fields() - 2) + - sizeof(int64_t) + sizeof(int32_t); - size_t rows = table.get()->num_rows(); - return Result({table, rows, rows * row_size}); + ARROW_ASSIGN_OR_RAISE( + table, + WithFixedSizeStringKeys(std::move(table), properties.num_ids, string_key_bytes)); + size_t rows = table->num_rows(); + size_t bytes = static_cast(util::TotalBufferSize(*table)); + return Result({std::move(table), rows, bytes}); } static void TableJoinOverhead(benchmark::State& state, @@ -54,10 +99,12 @@ static void TableJoinOverhead(benchmark::State& state, TableGenerationProperties right_table_properties, int batch_size, int num_right_tables, std::string factory_name, - std::shared_ptr options) { + std::shared_ptr options, bool use_threads, + int64_t string_key_bytes = 0) { left_table_properties.column_prefix = "lt"; left_table_properties.seed = 0; - ASSERT_OK_AND_ASSIGN(TableStats left_table_stats, MakeTable(left_table_properties)); + ASSERT_OK_AND_ASSIGN(TableStats left_table_stats, + MakeTable(left_table_properties, string_key_bytes)); size_t right_hand_rows = 0; size_t right_hand_bytes = 0; @@ -67,7 +114,8 @@ static void TableJoinOverhead(benchmark::State& state, for (int i = 0; i < num_right_tables; i++) { right_table_properties.column_prefix = "rt" + std::to_string(i); right_table_properties.seed = i + 1; - ASSERT_OK_AND_ASSIGN(TableStats right_table_stats, MakeTable(right_table_properties)); + ASSERT_OK_AND_ASSIGN(TableStats right_table_stats, + MakeTable(right_table_properties, string_key_bytes)); right_hand_rows += right_table_stats.rows; right_hand_bytes += right_table_stats.bytes; right_input_tables.push_back(std::move(right_table_stats)); @@ -86,9 +134,7 @@ static void TableJoinOverhead(benchmark::State& state, } Declaration join_node{factory_name, {input_nodes}, options}; state.ResumeTiming(); - // asof-join must currently be run synchronously as it relies on data arriving - // in-order - ASSERT_OK(DeclarationToStatus(std::move(join_node), /*use_threads=*/false)); + ASSERT_OK(DeclarationToStatus(std::move(join_node), use_threads)); } state.counters["rows_per_second"] = benchmark::Counter( @@ -113,7 +159,7 @@ AsofJoinNodeOptions GetRepeatedOptions(size_t repeat, FieldRef on_key, return AsofJoinNodeOptions(input_keys, tolerance); } -static void AsOfJoinOverhead(benchmark::State& state) { +static void AsOfJoinOverhead(benchmark::State& state, bool use_threads) { int64_t tolerance = 0; auto options = std::make_shared( GetRepeatedOptions(int(state.range(4) + 1), kTimeCol, {kKeyCol}, tolerance)); @@ -125,7 +171,28 @@ static void AsOfJoinOverhead(benchmark::State& state) { TableGenerationProperties{int(state.range(5)), int(state.range(6)), int(state.range(7)), "", kDefaultMinColumnVal, kDefaultMaxColumnVal, 0, kDefaultStart, kDefaultEnd}, - int(state.range(3)), int(state.range(4)), "asofjoin", std::move(options)); + int(state.range(3)), int(state.range(4)), "asofjoin", std::move(options), + use_threads); +} + +static void AsOfJoinKeyToleranceDensity(benchmark::State& state, bool use_threads, + int64_t string_key_bytes) { + constexpr int kColumns = 20; + constexpr int kIds = 500; + constexpr int kBatchSize = 4000; + constexpr int kNumRightTables = 1; + + auto options = std::make_shared( + GetRepeatedOptions(kNumRightTables + 1, kTimeCol, {kKeyCol}, state.range(2))); + TableJoinOverhead(state, + TableGenerationProperties{int(state.range(0)), kColumns, kIds, "", + kDefaultMinColumnVal, kDefaultMaxColumnVal, + 0, kDefaultStart, kDefaultEnd}, + TableGenerationProperties{int(state.range(1)), kColumns, kIds, "", + kDefaultMinColumnVal, kDefaultMaxColumnVal, + 0, kDefaultStart, kDefaultEnd}, + kBatchSize, kNumRightTables, "asofjoin", std::move(options), + use_threads, string_key_bytes); } // this generates the set of right hand tables to test on. @@ -163,7 +230,34 @@ void SetArgs(benchmark::internal::Benchmark* bench) { } } -BENCHMARK(AsOfJoinOverhead)->Apply(SetArgs); +void SetKeyToleranceDensityArgs(benchmark::internal::Benchmark* bench) { + bench->ArgNames({"left_freq", "right_freq", "tolerance"})->UseRealTime(); + + // A smaller frequency means a denser input. Include balanced inputs, dense left with + // sparse right (repeated right matches), and sparse left with dense right (many right + // candidates per left row). + for (const auto& [left_freq, right_freq] : + {std::pair{400, 400}, {200, 1000}, {1000, 200}}) { + for (int64_t tolerance : {-1000, 1000}) { + bench->Args({left_freq, right_freq, tolerance}); + } + } +} + +BENCHMARK_CAPTURE(AsOfJoinOverhead, serial_executor, false)->Apply(SetArgs); +BENCHMARK_CAPTURE(AsOfJoinKeyToleranceDensity, serial_executor_int32_keys, false, 0) + ->Apply(SetKeyToleranceDensityArgs); +BENCHMARK_CAPTURE(AsOfJoinKeyToleranceDensity, serial_executor_string128_keys, false, + kLongStringKeyBytes) + ->Apply(SetKeyToleranceDensityArgs); +#ifdef ARROW_ENABLE_THREADING +BENCHMARK_CAPTURE(AsOfJoinOverhead, threaded_executor, true)->Apply(SetArgs); +BENCHMARK_CAPTURE(AsOfJoinKeyToleranceDensity, threaded_executor_int32_keys, true, 0) + ->Apply(SetKeyToleranceDensityArgs); +BENCHMARK_CAPTURE(AsOfJoinKeyToleranceDensity, threaded_executor_string128_keys, true, + kLongStringKeyBytes) + ->Apply(SetKeyToleranceDensityArgs); +#endif } // namespace acero } // namespace arrow diff --git a/cpp/src/arrow/acero/asof_join_node.cc b/cpp/src/arrow/acero/asof_join_node.cc index 3970050e5022..947389ac661b 100644 --- a/cpp/src/arrow/acero/asof_join_node.cc +++ b/cpp/src/arrow/acero/asof_join_node.cc @@ -16,1174 +16,685 @@ // under the License. #include "arrow/acero/asof_join_node.h" -#include "arrow/acero/accumulation_queue.h" -#include "arrow/acero/backpressure_handler.h" -#include "arrow/acero/concurrent_queue_internal.h" +#include #include -#include +#include +#include +#include #include #include #include #include +#include #include -#include #include -#include +#include +#include +#include +#include "arrow/acero/accumulation_queue.h" #include "arrow/acero/exec_plan.h" #include "arrow/acero/exec_plan_internal.h" #include "arrow/acero/options.h" -#include "arrow/acero/unmaterialized_table_internal.h" -#ifndef NDEBUG -# include "arrow/acero/options_internal.h" -#endif #include "arrow/acero/query_context.h" -#include "arrow/acero/schema_util.h" -#include "arrow/acero/util.h" -#include "arrow/array/builder_binary.h" -#include "arrow/array/builder_primitive.h" -#ifndef NDEBUG -# include "arrow/compute/function_internal.h" -#endif #include "arrow/acero/time_series_util.h" +#include "arrow/array/builder_base.h" +#include "arrow/array/util.h" #include "arrow/compute/key_hash_internal.h" -#include "arrow/compute/light_array_internal.h" -#include "arrow/record_batch.h" #include "arrow/result.h" #include "arrow/status.h" #include "arrow/type_traits.h" -#include "arrow/util/bit_util.h" #include "arrow/util/checked_cast.h" -#include "arrow/util/config.h" -#include "arrow/util/future.h" #include "arrow/util/logging_internal.h" #include "arrow/util/string.h" namespace arrow { -using internal::ToChars; - -using compute::ColumnMetadataFromDataType; -using compute::Hashing64; -using compute::KeyColumnArray; -using compute::KeyColumnMetadata; -using compute::LightContext; +using compute::ExecBatch; +using compute::NullPlacement; using compute::SortKey; +using compute::SortOrder; +using internal::checked_cast; +using internal::ToChars; namespace acero { +namespace { -template -inline typename T::const_iterator std_find(const T& container, const V& val) { - return std::find(container.begin(), container.end(), val); -} - -template -inline bool std_has(const T& container, const V& val) { - return container.end() != std_find(container, val); -} - -template -inline D std_index(const T& container, const V& val) { - return std_find(container, val) - container.begin(); -} - -typedef uint64_t ByType; -typedef uint64_t OnType; -typedef uint64_t HashType; +using OnType = uint64_t; +using col_index_t = int; +using Task = util::SequencingQueue::Task; -/// A tolerance type with overflow-avoiding operations -struct TolType { - constexpr static OnType kMinValue = std::numeric_limits::lowest(); - constexpr static OnType kMaxValue = std::numeric_limits::max(); +enum class CandidateMode { Latest, Ordered }; - explicit TolType(int64_t tol) - : value(static_cast(tol > 0 ? tol : -tol)), positive(tol > 0) {} +class Tolerance { + public: + struct Bounds { + OnType lower; + OnType upper; + }; - OnType value; - bool positive; + explicit Tolerance(AsofJoinNodeOptions::ToleranceRange tolerance, + bool prefer_earlier_on_tie) + : lower_(tolerance.lower), + upper_(tolerance.upper), + prefer_earlier_on_tie_(prefer_earlier_on_tie) {} - // an entry with a time below this threshold expires - inline OnType Expiry(OnType left_value) { - return positive ? left_value - : (left_value < kMinValue + value ? kMinValue : left_value - value); + CandidateMode mode() const { + return upper_ <= 0 ? CandidateMode::Latest : CandidateMode::Ordered; } - // an entry with a time after this threshold is distant - inline OnType Horizon(OnType left_value) { - return positive ? (left_value > kMaxValue - value ? kMaxValue : left_value + value) - : left_value; - } + bool prefer_earlier_on_tie() const { return prefer_earlier_on_tie_; } - // true when the tolerance accepts the RHS time given the LHS one - inline bool Accepts(OnType left_value, OnType right_value) { - return positive - ? (left_value > right_value ? false : right_value - left_value <= value) - : (left_value < right_value ? false : left_value - right_value <= value); + std::optional BoundsFor(OnType left) const { + const OffsetResult lower = AddOffset(left, lower_); + const OffsetResult upper = AddOffset(left, upper_); + if (lower.above_max || upper.below_min) { + return std::nullopt; + } + DCHECK_LE(lower.value, upper.value); + return Bounds{lower.value, upper.value}; } -}; -// Maximum number of tables that can be joined -#define MAX_JOIN_TABLES 64 -typedef uint64_t row_index_t; -typedef int col_index_t; + private: + struct OffsetResult { + OnType value; + bool below_min = false; + bool above_max = false; + }; -// indicates normalization of a key value -template ::value, bool> = true> -static inline uint64_t key_value(T t) { - return static_cast(t); -} + static uint64_t Magnitude(int64_t offset) { + return offset >= 0 ? static_cast(offset) + : static_cast(-(offset + 1)) + uint64_t{1}; + } -class AsofJoinNode; + static OffsetResult AddOffset(OnType value, int64_t offset) { + const uint64_t magnitude = Magnitude(offset); + if (offset < 0) { + if (value < magnitude) { + return OffsetResult{0, /*below_min=*/true}; + } + return OffsetResult{value - magnitude}; + } + if (value > std::numeric_limits::max() - magnitude) { + return OffsetResult{std::numeric_limits::max(), /*below_min=*/false, + /*above_max=*/true}; + } + return OffsetResult{value + magnitude}; + } -#ifndef NDEBUG -// Get the debug-stream associated with the as-of-join node -std::ostream* GetDebugStream(AsofJoinNode* node); + const int64_t lower_; + const int64_t upper_; + const bool prefer_earlier_on_tie_; +}; -// Get the debug-mutex associated with the as-of-join node -std::mutex* GetDebugMutex(AsofJoinNode* node); +template +Result> ReadTimeValue(const Datum& value, int64_t row) { + using ArrowType = typename TypeIdTraits::Type; + using CType = typename TypeTraits::CType; + using ScalarType = typename TypeTraits::ScalarType; -// A debug-facility that wraps output-stream insertions with synchronization. Code like -// -// DebugSync(as_of_join_node) << ... << ... << ... ; -// -// will insert to the node's debug-stream and guard all insertions as one operation using -// the node's debug-mutex. -// -// However, it is recommended to use the DEBUG_SYNC macro, defined below it. Code like -// -// DEBUG_SYNC(as_of_join_node, ..., ..., ...); -// -// will do the same if NDEBUG is not defined, and otherwise it will be preprocessed-out. -// A IO manipulator used within DEBUG_SYNC must be wrapped by DEBUG_MANIP, for example: -// -// DEBUG_SYNC(as_of_join_node, ... , DEBUG_MANIP(std::endl) , ...); -class DebugSync { - public: - explicit DebugSync(AsofJoinNode* node) - : debug_os_(GetDebugStream(node)), - debug_mutex_(GetDebugMutex(node)), - alt_debug_mutex_(), // an alternative debug-mutex, if the node has none - debug_lock_(debug_mutex_ ? *debug_mutex_ : alt_debug_mutex_) { - if (debug_os_) { - std::ios state(NULL); - state.copyfmt(*debug_os_); - (*debug_os_) << "AsofjoinNode(" << std::hex << &node << "): "; - debug_os_->copyfmt(state); + if (value.is_scalar()) { + const auto& scalar = checked_cast(*value.scalar()); + if (!scalar.is_valid) { + return std::nullopt; } + return NormalizeTime(static_cast(scalar.value)); } - - DebugSync& operator<<(std::ostream& (*pf)(std::ostream&)) { - if (debug_os_) pf(*debug_os_); - return *this; - } - DebugSync& operator<<(std::ios& (*pf)(std::ios&)) { - if (debug_os_) pf(*debug_os_); - return *this; - } - DebugSync& operator<<(std::ios_base& (*pf)(std::ios_base&)) { - if (debug_os_) pf(*debug_os_); - return *this; + if (value.is_array()) { + ArraySpan array(*value.array()); + if (array.IsNull(row)) { + return std::nullopt; + } + return NormalizeTime(array.GetValues(1)[row]); } + return Status::Invalid("AsofJoin on-key must be an array or scalar, but got ", + ::arrow::ToString(value.kind())); +} - // used by DEBUG_MANIP macro below - using Manip = std::function; - DebugSync& operator<<(Manip f) { return f(*this); } - - template - DebugSync& operator<<(T&& value) { - if (debug_os_) (*debug_os_) << value; - return *this; +Result> ReadTimeValue(const Datum& value, int64_t row) { + switch (value.type()->id()) { +#define ASOF_TIME_CASE(ID) \ + case Type::ID: \ + return ReadTimeValue(value, row) + ASOF_TIME_CASE(INT8); + ASOF_TIME_CASE(INT16); + ASOF_TIME_CASE(INT32); + ASOF_TIME_CASE(INT64); + ASOF_TIME_CASE(UINT8); + ASOF_TIME_CASE(UINT16); + ASOF_TIME_CASE(UINT32); + ASOF_TIME_CASE(UINT64); + ASOF_TIME_CASE(DATE32); + ASOF_TIME_CASE(DATE64); + ASOF_TIME_CASE(TIME32); + ASOF_TIME_CASE(TIME64); + ASOF_TIME_CASE(TIMESTAMP); +#undef ASOF_TIME_CASE + default: + return Status::Invalid("Unsupported AsofJoin on-key type ", + value.type()->ToString()); } +} - // used by DEBUG_SYNC macro below - template - DebugSync& insert(Args&&... args) { - return (*this << ... << args); - } +struct PreparedKeyColumn { + int64_t Row(int64_t row) const { return is_scalar ? 0 : row; } - private: - std::ostream* debug_os_; - std::mutex* debug_mutex_; - std::mutex alt_debug_mutex_; - std::unique_lock debug_lock_; + std::shared_ptr values; + bool is_scalar; }; -# define DEBUG_SYNC(node, ...) DebugSync(node).insert(__VA_ARGS__) -# define DEBUG_MANIP(manip) \ - DebugSync::Manip([](DebugSync& d) -> DebugSync& { return d << manip; }) -# define NDEBUG_EXPLICIT -# define DEBUG_ADD(ndebug, ...) ndebug, __VA_ARGS__ -#else -# define DEBUG_SYNC(...) -# define DEBUG_MANIP(...) -# define NDEBUG_EXPLICIT explicit -# define DEBUG_ADD(ndebug, ...) ndebug -#endif - -struct MemoStore { - // A MemoStore is associated with an input of the as-of-join node. - // Stores the current time as well as the last-known values for all the keys. - // In case of a future as-of-join, also stores: - // 1. known future values for all keys - // 2. known distinct future times - // - // Keeping future values and times allows the as-of-join node's to look ahead to the - // horizon (i.e., the left input's time plus future tolerance) of a right input. - // - // A key's value is captured in an entry, which describes a row of the input's batch. - - struct Entry { - Entry() = default; - - Entry(OnType time, std::shared_ptr batch, row_index_t row) - : time(time), batch(batch), row(row) {} - - void swap(Entry& other) { - std::swap(time, other.time); - std::swap(batch, other.batch); - std::swap(row, other.row); - } - - // Timestamp associated with the entry - OnType time; - - // Batch associated with the entry (perf is probably OK for this; batches change - // rarely) - std::shared_ptr batch; - - // Row associated with the entry - row_index_t row; - }; - - NDEBUG_EXPLICIT MemoStore(DEBUG_ADD(bool no_future, AsofJoinNode* node, size_t index)) - : no_future_(no_future), - DEBUG_ADD(current_time_(std::numeric_limits::lowest()), node_(node), - index_(index)) {} - - // true when there are no future entries, which is the case for the LHS table and the - // case for when the tolerance is non-positive. A non-positive-tolerance as-of-join - // operation requires memorizing only the most recently observed entry per key. OTOH, a - // positive-tolerance (future) as-of-join operation requires memorizing per-key queues - // of entries up to the tolerance's horizon and in particular distinguishes between the - // current (front-of-queue) and latest (back-of-queue) entries per key. - bool no_future_; - // the time of the current entry, defaulting to 0. - // when entries with a time less than T are removed, the current time is updated to the - // time of the next (by-time) and now-current entry or to T if no such entry exists. - OnType current_time_; - // current entry per key - std::unordered_map entries_; - // future entries per key - std::unordered_map> future_entries_; - // current and future (distinct) times of existing entries - std::deque times_; -#ifndef NDEBUG - // Owning node - AsofJoinNode* node_; - // Index of owning input - size_t index_; -#endif - - void swap(MemoStore& memo) { -#ifndef NDEBUG - std::swap(node_, memo.node_); - std::swap(index_, memo.index_); -#endif - std::swap(no_future_, memo.no_future_); - std::swap(current_time_, memo.current_time_); - entries_.swap(memo.entries_); - future_entries_.swap(memo.future_entries_); - times_.swap(memo.times_); - } - - // Updates the current time to `ts` if it is less. Returns true if updated. - bool UpdateTime(OnType ts) { - bool update = current_time_ < ts; - if (update) current_time_ = ts; - return update; - } - - void Store(const std::shared_ptr& batch, row_index_t row, OnType time, - DEBUG_ADD(ByType key, OnType for_time)) { - DEBUG_SYNC(node_, "memo ", index_, " store: for_time=", for_time, " row=", row, - " time=", time, " key=", key, DEBUG_MANIP(std::endl)); - if (no_future_ || entries_.count(key) == 0) { - auto& e = entries_[key]; - // that we can do this assignment optionally, is why we - // can get away with using shared_ptr above (the batch - // shouldn't change that often) - if (e.batch != batch) e.batch = batch; - e.row = row; - e.time = time; - } else { - future_entries_[key].emplace(time, batch, row); - } - // Maintain distinct times: - // If no times are currently maintained then the given time is distinct and hence - // pushed. Otherwise, the invariant is that the latest time is at the back. If no - // future times are maintained, then only one time is maintained, and hence it is - // overwritten with the given time. Otherwise, the given time must be no less than the - // latest time, due to time ordering, so it is pushed back only if it is distinct. - if (times_.empty() || (!no_future_ && times_.back() != time)) { - times_.push_back(time); - } else { - times_.back() = time; +struct PreparedBatch { + uint64_t Hash(int64_t row) const { return hashes.empty() ? 0 : hashes[row]; } + + bool KeysEqual(int64_t row, const PreparedBatch& other, int64_t other_row) const { + DCHECK_EQ(key_columns.size(), other.key_columns.size()); + for (size_t i = 0; i < key_columns.size(); ++i) { + const PreparedKeyColumn& left = key_columns[i]; + const PreparedKeyColumn& right = other.key_columns[i]; + const int64_t left_row = left.Row(row); + const int64_t right_row = right.Row(other_row); + if (!left.values->RangeEquals(*right.values, left_row, left_row + 1, right_row)) { + return false; + } } - // `time` is the most advanced seen yet - `UpdateTime(time)` would work but not needed - current_time_ = time; + return true; } - std::optional GetEntryForKey(ByType key) const { - auto e = entries_.find(key); - return entries_.end() == e ? std::nullopt : std::optional(&e->second); - } + ExecBatch batch; + std::vector> times; + std::vector key_columns; + std::vector hashes; +}; - bool RemoveEntriesWithLesserTime(OnType ts) { - DEBUG_SYNC(node_, "memo ", index_, " remove: ts=", ts, DEBUG_MANIP(std::endl)); - bool updated = false; - // remove future entries with lesser time - for (auto fe = future_entries_.begin(); fe != future_entries_.end();) { - auto& queue = fe->second; - while (!queue.empty() && queue.front().time < ts) { - queue.pop(); - updated = true; // queue changed - } - // remove entry if its queue was just emptied - if (queue.empty()) { - fe = future_entries_.erase(fe); +Result> PrepareBatch( + ExecBatch batch, col_index_t on_key, const std::vector& by_keys, + compute::ExecContext* ctx) { + if (batch.length > std::numeric_limits::max()) { + return Status::CapacityError("AsofJoin input batch has too many rows: ", + batch.length); + } + + auto prepared = std::make_shared(); + prepared->batch = std::move(batch); + prepared->times.reserve(prepared->batch.length); + for (int64_t row = 0; row < prepared->batch.length; ++row) { + ARROW_ASSIGN_OR_RAISE(auto time, ReadTimeValue(prepared->batch.values[on_key], row)); + prepared->times.push_back(time); + } + + if (!by_keys.empty()) { + std::vector key_values; + key_values.reserve(by_keys.size()); + prepared->key_columns.reserve(by_keys.size()); + for (col_index_t by_key : by_keys) { + const Datum& value = prepared->batch.values[by_key]; + if (value.is_scalar()) { + ARROW_ASSIGN_OR_RAISE( + auto scalar_array, + MakeArrayFromScalar(*value.scalar(), /*length=*/1, ctx->memory_pool())); + prepared->key_columns.push_back({std::move(scalar_array), true}); + ARROW_ASSIGN_OR_RAISE(auto hash_array, + MakeArrayFromScalar(*value.scalar(), prepared->batch.length, + ctx->memory_pool())); + key_values.emplace_back(std::move(hash_array)); } else { - ++fe; + key_values.push_back(value); + prepared->key_columns.push_back({MakeArray(value.array()), false}); } } - // remove last-known entries with lesser time - for (auto e = entries_.begin(); e != entries_.end();) { - if (e->second.time < ts) { - // drop last-known entry and move next future entry, if exists, in its place - auto fe = future_entries_.find(e->first); - if (fe != future_entries_.end() && !fe->second.empty()) { - auto& queue = fe->second; - e->second.swap(queue.front()); - queue.pop(); - ++e; - } else { - e = entries_.erase(e); - } - updated = true; // entry changed - } else { - ++e; - } - } - // remove known lesser times - while (!times_.empty() && times_.front() < ts) { - times_.pop_front(); - } - // update current time - return UpdateTime(ts) || updated; - } -}; - -// a specialized higher-performance variation of Hashing64 logic from hash_join_node -// the code here avoids recreating objects that are independent of each batch processed -class KeyHasher { - friend class AsofJoinNode; - - static constexpr int kMiniBatchLength = arrow::util::MiniBatch::kMiniBatchLength; - public: - // the key hasher is not thread-safe and is only used in sequential batch processing - // of the input it is associated with - KeyHasher(size_t index, const std::vector& indices) - : index_(index), - indices_(indices), - metadata_(indices.size()), - batch_(NULLPTR), - hashes_(), - ctx_(), - column_arrays_(), - stack_() { - ctx_.stack = &stack_; - column_arrays_.resize(indices.size()); - } - - Status Init(ExecContext* exec_context, const std::shared_ptr& schema) { - ctx_.hardware_flags = exec_context->cpu_info()->hardware_flags(); - const auto& fields = schema->fields(); - for (size_t k = 0; k < metadata_.size(); k++) { - ARROW_ASSIGN_OR_RAISE(metadata_[k], - ColumnMetadataFromDataType(fields[indices_[k]]->type())); - } - return stack_.Init(exec_context->memory_pool(), - 4 * kMiniBatchLength * sizeof(uint32_t)); - } - - // invalidate cached hashes for batch - required when it changes - // only this method can be called concurrently with HashesFor - void Invalidate() { batch_ = NULLPTR; } - - // compute and cache a hash for each row of the given batch - const std::vector& HashesFor(const RecordBatch* batch) { - if (batch_ == batch) { - return hashes_; // cache hit - return cached hashes - } - Invalidate(); - size_t batch_length = batch->num_rows(); - hashes_.resize(batch_length); - for (int64_t i = 0; i < static_cast(batch_length); i += kMiniBatchLength) { - int64_t length = std::min(static_cast(batch_length - i), - static_cast(kMiniBatchLength)); - for (size_t k = 0; k < indices_.size(); k++) { - auto array_data = batch->column_data(indices_[k]); - column_arrays_[k] = - ColumnArrayFromArrayDataAndMetadata(array_data, metadata_[k], i, length); - } - // write directly to the cache - Hashing64::HashMultiColumn(column_arrays_, &ctx_, hashes_.data() + i); + prepared->hashes.resize(prepared->batch.length); + if (prepared->batch.length > 0) { + ExecBatch key_batch(std::move(key_values), prepared->batch.length); + ::arrow::util::TempVectorStack temp_stack; + ARROW_RETURN_NOT_OK(temp_stack.Init(ctx->memory_pool(), + compute::Hashing64::kHashBatchTempStackUsage)); + std::vector column_arrays; + ARROW_RETURN_NOT_OK( + compute::Hashing64::HashBatch(key_batch, prepared->hashes.data(), column_arrays, + ctx->cpu_info()->hardware_flags(), &temp_stack, + /*start_row=*/0, prepared->batch.length)); } - DEBUG_SYNC(node_, "key hasher ", index_, " got hashes ", - compute::internal::GenericToString(hashes_), DEBUG_MANIP(std::endl)); - batch_ = batch; // associate cache with current batch - return hashes_; } + return prepared; +} - private: - AsofJoinNode* node_ = nullptr; // avoids circular dependency during initialization - size_t index_; - std::vector indices_; - std::vector metadata_; - std::atomic batch_; - std::vector hashes_; - LightContext ctx_; - std::vector column_arrays_; - arrow::util::TempVectorStack stack_; +struct RowRef { + std::shared_ptr batch; + int64_t row; }; -class BackpressureController : public BackpressureControl { - public: - BackpressureController(ExecNode* node, ExecNode* output, - std::atomic& backpressure_counter) - : node_(node), output_(output), backpressure_counter_(backpressure_counter) {} - - void Pause() override { node_->PauseProducing(output_, ++backpressure_counter_); } - void Resume() override { node_->ResumeProducing(output_, ++backpressure_counter_); } - - private: - ExecNode* node_; - ExecNode* output_; - std::atomic& backpressure_counter_; -}; +bool KeysEqual(const RowRef& left, const RowRef& right) { + return left.batch->KeysEqual(left.row, *right.batch, right.row); +} -class InputState : public util::SerialSequencingQueue::Processor { - // InputState corresponds to an input - // Input record batches are queued up in InputState until processed and - // turned into output record batches. +struct FlowAction { + enum class Kind { None, Pause, Resume }; - public: - InputState(size_t index, TolType tolerance, bool must_hash, bool may_rehash, - KeyHasher* key_hasher, AsofJoinNode* node, BackpressureHandler handler, - const std::shared_ptr& schema, - const col_index_t time_col_index, - const std::vector& key_col_index) - : sequencer_(util::SerialSequencingQueue::Make(this)), - queue_(std::move(handler)), - schema_(schema), - time_col_index_(time_col_index), - key_col_index_(key_col_index), - time_type_id_(schema_->fields()[time_col_index_]->type()->id()), - key_type_id_(key_col_index.size()), - key_hasher_(key_hasher), - node_(node), - index_(index), - must_hash_(must_hash), - may_rehash_(may_rehash), - tolerance_(tolerance), - memo_(DEBUG_ADD(/*no_future=*/index == 0 || !tolerance.positive, node, index)) { - for (size_t k = 0; k < key_col_index_.size(); k++) { - key_type_id_[k] = schema_->fields()[key_col_index_[k]]->type()->id(); + void Apply() const { + if (kind == Kind::Pause) { + input->PauseProducing(output, counter); + } else if (kind == Kind::Resume) { + input->ResumeProducing(output, counter); } } - static Result> Make( - size_t index, TolType tolerance, bool must_hash, bool may_rehash, - KeyHasher* key_hasher, ExecNode* asof_input, AsofJoinNode* asof_node, - std::atomic& backpressure_counter, - const std::shared_ptr& schema, const col_index_t time_col_index, - const std::vector& key_col_index) { - constexpr size_t low_threshold = 4, high_threshold = 8; - std::unique_ptr backpressure_control = - std::make_unique( - /*node=*/asof_input, /*output=*/asof_node, backpressure_counter); - ARROW_ASSIGN_OR_RAISE(auto handler, - BackpressureHandler::Make(low_threshold, high_threshold, - std::move(backpressure_control))); - return std::make_unique(index, tolerance, must_hash, may_rehash, - key_hasher, asof_node, std::move(handler), schema, - time_col_index, key_col_index); - } - - col_index_t InitSrcToDstMapping(col_index_t dst_offset, bool skip_time_and_key_fields) { - src_to_dst_.resize(schema_->num_fields()); - for (int i = 0; i < schema_->num_fields(); ++i) - if (!(skip_time_and_key_fields && IsTimeOrKeyColumn(i))) - src_to_dst_[i] = dst_offset++; - return dst_offset; - } - - const std::optional& MapSrcToDst(col_index_t src) const { - return src_to_dst_[src]; - } - - bool IsTimeOrKeyColumn(col_index_t i) const { - DCHECK_LT(i, schema_->num_fields()); - return (i == time_col_index_) || std_has(key_col_index_, i); - } + Kind kind = Kind::None; + ExecNode* input = nullptr; + ExecNode* output = nullptr; + int32_t counter = 0; +}; - // Gets the latest row index, assuming the queue isn't empty - row_index_t GetLatestRow() const { return latest_ref_row_; } +class AsofJoinNode; - bool Empty() const { - // cannot be empty if ref row is >0 -- can avoid slow queue lock - // below - if (latest_ref_row_ > 0) return false; - return queue_.Empty(); - } +class InputState final : public util::SerialSequencingQueue::Processor { + public: + InputState(AsofJoinNode* node, size_t index, ExecNode* input, col_index_t on_key, + std::vector by_keys, + std::optional null_placement); - // 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. - // 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); - } + Status InsertBatch(ExecBatch batch); + Status Process(ExecBatch batch) override; - // in case memo may not have future entries (the case of a non-positive tolerance), - // returns the latest time (which is current); otherwise, returns the current time. - // used when checking whether RHS is up to date with LHS. - OnType GetCurrentTime() const { - return memo_.no_future_ ? GetLatestTime() : static_cast(memo_.current_time_); - } + FlowAction BatchBuffered(); + FlowAction BatchConsumed(); + void Shutdown(); - int total_batches() const { return total_batches_; } + private: + FlowAction SetUpstreamPausedUnlocked(bool paused); + Status ValidateTimes(const PreparedBatch& batch); - // Gets latest batch (precondition: must not be empty) - const std::shared_ptr& GetLatestBatch() const { - return queue_.Front(); - } + static constexpr size_t kLowWatermark = 4; + static constexpr size_t kHighWatermark = 8; -#define LATEST_VAL_CASE(id, val) \ - case Type::id: { \ - using T = typename TypeIdTraits::Type; \ - using CType = typename TypeTraits::CType; \ - return val(data->GetValues(1)[row]); \ - } + AsofJoinNode* node_; + size_t index_; + ExecNode* input_; + col_index_t on_key_; + std::vector by_keys_; + std::unique_ptr sequencer_; - inline ByType GetLatestKey() const { - return GetKey(GetLatestBatch().get(), latest_ref_row_); - } + std::optional last_time_; + std::optional null_placement_; + bool saw_trailing_null_ = false; - inline ByType GetKey(const RecordBatch* batch, row_index_t row) const { - if (must_hash_) { - // Query the key hasher. This may hit cache, which must be valid for the batch. - // Therefore, the key hasher is invalidated when a new batch is pushed - see - // `InputState::Push`. - return key_hasher_->HashesFor(batch)[row]; - } - if (key_col_index_.size() == 0) { - return 0; - } - auto data = batch->column_data(key_col_index_[0]); - switch (key_type_id_[0]) { - LATEST_VAL_CASE(INT8, key_value) - LATEST_VAL_CASE(INT16, key_value) - LATEST_VAL_CASE(INT32, key_value) - LATEST_VAL_CASE(INT64, key_value) - LATEST_VAL_CASE(UINT8, key_value) - LATEST_VAL_CASE(UINT16, key_value) - LATEST_VAL_CASE(UINT32, key_value) - LATEST_VAL_CASE(UINT64, key_value) - LATEST_VAL_CASE(DATE32, key_value) - LATEST_VAL_CASE(DATE64, key_value) - LATEST_VAL_CASE(TIME32, key_value) - LATEST_VAL_CASE(TIME64, key_value) - LATEST_VAL_CASE(TIMESTAMP, key_value) - default: - DCHECK(false); - return 0; // cannot happen - } - } + std::mutex flow_mutex_; + size_t buffered_batches_ = 0; + bool upstream_paused_ = false; + int32_t outgoing_counter_ = 0; + bool shutdown_ = false; +}; - inline OnType GetLatestTime() const { - return GetTime(GetLatestBatch().get(), time_type_id_, time_col_index_, - latest_ref_row_); - } +class RhsLane { + public: + RhsLane(AsofJoinNode* node, size_t lane_index, size_t input_index, + std::vector payload_columns, Tolerance tolerance, MemoryPool* pool) + : node_(node), + lane_index_(lane_index), + input_index_(input_index), + payload_columns_(std::move(payload_columns)), + tolerance_(tolerance), + pool_(pool) {} -#undef LATEST_VAL_CASE + Result> Enqueue(std::shared_ptr batch); + Result> SetTotal(int total_batches); + Result Assign(std::shared_ptr left); + void Stop(); - bool Finished() const { return batches_processed_ == total_batches_; } + private: + enum class Phase { + // No left batch is assigned. Assign() installs a job, posts Run(), and moves to + // Claimed; Stop() moves to Stopped. + NoJob, + + // The current job reached the end of available RHS data. It waits for Enqueue() + // or RHS completion; the waking producer posts Run() and moves to Claimed. + // Stop() moves to Stopped. + Waiting, + + // One posted or running Run() task exclusively owns the job, preventing duplicate + // runners and lost wakes. Blocking moves to Waiting, completing the job moves to + // NoJob, and Stop() moves to Stopped. + Claimed, + + // The lane is terminal. It ignores new data and assignments and never exits. + Stopped, + }; - Result Advance() { - // Try advancing to the next row and update latest_ref_row_ - // Returns true if able to advance, false if not. - bool have_active_batch = - (latest_ref_row_ > 0 /*short circuit the lock on the queue*/) || !queue_.Empty(); + struct SelectionRun { + std::optional source; + int64_t length; + }; - if (have_active_batch) { - OnType next_time = GetLatestTime(); - if (latest_time_ > next_time) { - return Status::Invalid("AsofJoin does not allow out-of-order on-key values"); - } - latest_time_ = next_time; - // If we have an active batch - if (++latest_ref_row_ >= (row_index_t)queue_.Front()->num_rows()) { - // hit the end of the batch, need to get the next batch if possible. - ++batches_processed_; - latest_ref_row_ = 0; - bool did_pop = queue_.TryPop().has_value(); - DCHECK(did_pop); - ARROW_UNUSED(did_pop); - have_active_batch = !queue_.Empty(); - } - } - return have_active_batch; - } - - // Advance the data to be immediately past the tolerance's horizon for the specified - // 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. - // 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 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 - return memo_.no_future_ ? false : memo_.RemoveEntriesWithLesserTime(ts); - } - - // Not updated. Try to update and possibly advance. - bool advanced, updated = false; - OnType latest_time; - do { - latest_time = GetLatestTime(); - // if Advance() returns true, then the latest_ts must also be valid - // Keep advancing right table until we hit the latest row that has - // timestamp <= ts. This is because we only need the latest row for the - // match given a left ts. - if (latest_time > tolerance_.Horizon(ts)) { // hit a distant timestamp - DEBUG_SYNC(node_, "Advancing input ", index_, " hit distant time=", latest_time, - " at=", ts, DEBUG_MANIP(std::endl)); - // if no future entries, which would have been earlier than the distant time, no - // need to queue it - if (memo_.future_entries_.empty()) break; - } - auto rb = GetLatestBatch(); - if (may_rehash_ && rb->column_data(key_col_index_[0])->GetNullCount() > 0) { - must_hash_ = true; - may_rehash_ = false; - Rehash(); - } - memo_.Store(rb, latest_ref_row_, latest_time, DEBUG_ADD(GetLatestKey(), ts)); - // negative tolerance means a last-known entry was stored - set `updated` to `true` - updated = memo_.no_future_; - ARROW_ASSIGN_OR_RAISE(advanced, Advance()); - } while (advanced); - if (!memo_.no_future_ && latest_time >= ts) { - // `updated` was not modified in the loop from the initial `false` value; set it now - updated = memo_.RemoveEntriesWithLesserTime(ts); - } - DEBUG_SYNC(node_, "Advancing input ", index_, " updated=", updated, - DEBUG_MANIP(std::endl)); - return updated; - } - Status InsertBatch(ExecBatch batch) { - return sequencer_->InsertBatch(std::move(batch)); - } - - Status Process(ExecBatch batch) override { - auto rb = *batch.ToRecordBatch(schema_); - DEBUG_SYNC(node_, "received batch from input ", index_, ":", DEBUG_MANIP(std::endl), - rb->ToString(), DEBUG_MANIP(std::endl)); - return Push(rb); - } - void Rehash() { - DEBUG_SYNC(node_, "rehashing for input ", index_, ":", DEBUG_MANIP(std::endl)); - MemoStore new_memo(DEBUG_ADD(memo_.no_future_, node_, index_)); - new_memo.current_time_ = (OnType)memo_.current_time_; - for (auto e = memo_.entries_.begin(); e != memo_.entries_.end(); ++e) { - auto& entry = e->second; - auto new_key = GetKey(entry.batch.get(), entry.row); - DEBUG_SYNC(node_, " ", e->first, " to ", new_key, DEBUG_MANIP(std::endl)); - new_memo.entries_[new_key].swap(entry); - auto fe = memo_.future_entries_.find(e->first); - if (fe != memo_.future_entries_.end()) { - new_memo.future_entries_[new_key].swap(fe->second); + struct Job { + explicit Job(std::shared_ptr left) : left(std::move(left)) {} + + void AppendMatch(std::optional match) { + if (!selections.empty()) { + SelectionRun& previous = selections.back(); + if ((!previous.source && !match) || + (previous.source && match && previous.source->batch == match->batch && + previous.source->row + previous.length == match->row)) { + ++previous.length; + return; + } } + selections.push_back({std::move(match), 1}); } - memo_.times_.swap(new_memo.times_); - memo_.swap(new_memo); - } - Status Push(const std::shared_ptr& rb) { - if (rb->num_rows() > 0) { - key_hasher_->Invalidate(); // batch changed - invalidate key hasher's cache - queue_.Push(rb); // only now push batch for processing - } else { - ++batches_processed_; // don't enqueue empty batches, just record as processed - } - return Status::OK(); - } + std::shared_ptr left; + int64_t left_row = 0; + std::vector selections; + }; - std::optional GetMemoEntryForKey(ByType key) { - return memo_.GetEntryForKey(key); - } + struct Candidate { + RowRef row; + uint64_t version; + }; - std::optional GetMemoTimeForKey(ByType key) { - auto r = GetMemoEntryForKey(key); - if (r.has_value()) { - return (*r)->time; - } else { - return std::nullopt; - } - } + struct OrderedCandidate { + RowRef row; + OnType time; + uint64_t version; + }; - void RemoveMemoEntriesWithLesserTime(OnType ts) { - memo_.RemoveEntriesWithLesserTime(ts); - } + struct OrderedCandidates { + std::deque rows; + }; - const std::shared_ptr& get_schema() const { return schema_; } + struct ExpiryEntry { + OnType time; + uint64_t hash; + uint64_t version; + }; - void set_total_batches(int n) { - DCHECK_GE(n, 0); - DCHECK_EQ(total_batches_, -1) << "Set total batch more than once"; - total_batches_ = n; - } + struct PeekResult { + enum class Kind { Row, Blocked, End }; + Kind kind; + RowRef row; + }; - void ForceShutdown() { - // Force the upstream input node to unpause. Necessary to avoid deadlock when we - // terminate the process thread - queue_.ForceShutdown(); - } + Status Run(); + Result PeekNext(); + void ConsumeNext(); + bool StreamEndedUnlocked() const; + bool WaitOrRetry(); + void RememberBackward(const RowRef& row, OnType time); + void ExpireBackward(OnType lower_bound); + void RememberOrdered(const RowRef& row, OnType time); + void ExpireOrdered(OnType lower_bound); + std::optional MatchBackward(const RowRef& key) const; + std::optional MatchOrdered(const RowRef& key, OnType left_time) const; + Result> Materialize(const Job& job) const; - private: - // ExecBatch Sequencer - std::unique_ptr sequencer_; - // Pending record batches. The latest is the front. Batches cannot be empty. - BackpressureConcurrentQueue> queue_; - // Schema associated with the input - std::shared_ptr schema_; - // Total number of batches (only int because InputFinished uses int) - std::atomic total_batches_{-1}; - // Number of batches processed so far (only int because InputFinished uses int) - std::atomic batches_processed_{0}; - // Index of the time col - col_index_t time_col_index_; - // Index of the key col - std::vector key_col_index_; - // Type id of the time column - Type::type time_type_id_; - // Type id of the key column - std::vector key_type_id_; - // Hasher for key elements - mutable KeyHasher* key_hasher_; - // Owning node AsofJoinNode* node_; - // Index of this input - size_t index_; - // True if hashing is mandatory - bool must_hash_; - // True if by-key values may be rehashed - bool may_rehash_; - // Tolerance - TolType tolerance_; - // Index of the latest row reference within; if >0 then queue_ cannot be empty - // Must be < queue_.front()->num_rows() if queue_ is non-empty - row_index_t latest_ref_row_ = 0; - // Time of latest row - OnType latest_time_ = std::numeric_limits::lowest(); - // Stores latest known values for the various keys - MemoStore memo_; - // Mapping of source columns to destination columns - std::vector> src_to_dst_; + size_t lane_index_; + size_t input_index_; + std::vector payload_columns_; + Tolerance tolerance_; + MemoryPool* pool_; + + std::mutex mutex_; + Phase phase_ = Phase::NoJob; + std::shared_ptr job_; + std::deque> batches_; + // Only the lane owner touches the active batch and row. Producers append to + // batches_ under mutex_ and ownership transfers between lane tasks through phase_. + std::shared_ptr current_batch_; + int64_t current_row_ = 0; + int received_batches_ = 0; + std::optional total_batches_; + + uint64_t next_version_ = 0; + std::unordered_map> backward_candidates_; + std::deque backward_expiry_; + std::unordered_map> ordered_candidates_; + std::deque ordered_expiry_; }; -/// Wrapper around UnmaterializedCompositeTable that knows how to emplace -/// the join row-by-row -template -class CompositeTableBuilder { - using SliceBuilder = UnmaterializedSliceBuilder; - using CompositeTable = UnmaterializedCompositeTable; - - public: - NDEBUG_EXPLICIT CompositeTableBuilder( - const std::vector>& inputs, - const std::shared_ptr& schema, arrow::MemoryPool* pool, - DEBUG_ADD(size_t n_tables, AsofJoinNode* node)) - : unmaterialized_table(InitUnmaterializedTable(schema, inputs, pool)), - DEBUG_ADD(n_tables_(n_tables), node_(node)) { - DCHECK_GE(n_tables_, 1); - DCHECK_LE(n_tables_, MAX_TABLES); - } - - size_t n_rows() const { return unmaterialized_table.Size(); } - - // Adds the latest row from the input state as a new composite reference row - // - LHS must have a valid key,timestep,and latest rows - // - RHS must have valid data memo'ed for the key - void Emplace(std::vector>& in, TolType tolerance) { - DCHECK_EQ(in.size(), n_tables_); - - // Get the LHS key - ByType key = in[0]->GetLatestKey(); - - // Add row and setup LHS - // (the LHS state comes just from the latest row of the LHS table) - DCHECK(!in[0]->Empty()); - const std::shared_ptr& lhs_latest_batch = in[0]->GetLatestBatch(); - row_index_t lhs_latest_row = in[0]->GetLatestRow(); - OnType lhs_latest_time = in[0]->GetLatestTime(); - if (0 == lhs_latest_row) { - // On the first row of the batch, we resize the destination. - // The destination size is dictated by the size of the LHS batch. - row_index_t new_batch_size = lhs_latest_batch->num_rows(); - row_index_t new_capacity = unmaterialized_table.Size() + new_batch_size; - if (unmaterialized_table.capacity() < new_capacity) { - unmaterialized_table.reserve(new_capacity); - } - } - - SliceBuilder new_row{&unmaterialized_table}; +// No left batch is active. It waits for an activatable queued batch and then moves to +// Matching; exhausting the left input with an empty queue, or Stop(), moves to Terminal. +struct WaitingForLeft {}; + +// Owns the active left batch while RHS lanes match it in parallel. It waits for every +// lane result; the last result assembles the output and moves to OutputInFlight. Stop() +// moves to Terminal. +struct Matching { + explicit Matching(std::shared_ptr left, size_t lane_count) + : left(std::move(left)), results(lane_count), remaining(lane_count) {} + std::shared_ptr left; + std::vector>> results; + size_t remaining; +}; - // Each item represents a portion of the columns of the output table - new_row.AddEntry(lhs_latest_batch, lhs_latest_row, lhs_latest_row + 1); +// Keeps the active left batch alive while its assembled output is delivered. When +// downstream InputReceived succeeds, output completion moves to WaitingForLeft. Stop() +// moves to Terminal. +struct OutputInFlight { + explicit OutputInFlight(std::shared_ptr left) : left(std::move(left)) {} + std::shared_ptr left; +}; - DEBUG_SYNC(node_, "Emplace: key=", key, " lhs_latest_row=", lhs_latest_row, - " lhs_latest_time=", lhs_latest_time, DEBUG_MANIP(std::endl)); +// No coordinator work remains. Normal left-input exhaustion or Stop() enters this +// state; subsequent coordinator events are ignored and it never exits. +struct Terminal {}; - // Get the state for that key from all on the RHS -- assumes it's up to date - // (the RHS state comes from the memoized row references) - for (size_t i = 1; i < in.size(); ++i) { - std::optional opt_entry = in[i]->GetMemoEntryForKey(key); -#ifndef NDEBUG - { - bool has_entry = opt_entry.has_value(); - OnType entry_time = has_entry ? (*opt_entry)->time : TolType::kMinValue; - row_index_t entry_row = has_entry ? (*opt_entry)->row : 0; - bool accepted = has_entry && tolerance.Accepts(lhs_latest_time, entry_time); - DEBUG_SYNC(node_, " i=", i, " has_entry=", has_entry, " time=", entry_time, - " row=", entry_row, " accepted=", accepted, DEBUG_MANIP(std::endl)); - } -#endif - if (opt_entry.has_value()) { - DCHECK(*opt_entry); - if (tolerance.Accepts(lhs_latest_time, (*opt_entry)->time)) { - // Have a valid entry - const MemoStore::Entry* entry = *opt_entry; - new_row.AddEntry(entry->batch, entry->row, entry->row + 1); - continue; - } - } - new_row.AddEntry(nullptr, 0, 1); - } - new_row.Finalize(); - } +// Controls whether WaitingForLeft may activate another batch. +enum class LeftGate { + // Downstream accepts output, so a queued left batch may activate immediately. + // PauseProducing moves to Paused; sequencing the complete left input moves to Flushing. + Open, - // Materializes the current reference table into a target record batch - Result>> Materialize() { - return unmaterialized_table.Materialize(); - } + // Downstream backpressure holds queued left batches at the generation boundary. It + // waits for ResumeProducing and returns to Open, unless left completion moves it to + // Flushing first. + Paused, - // Returns true if there are no rows - bool empty() const { return unmaterialized_table.Empty(); } + // Every declared left batch is sequenced, so the fixed tail drains without waiting for + // downstream resume. RHS lanes may still wait for right data; this gate never exits. + Flushing, +}; - private: - CompositeTable unmaterialized_table; +using CoordinatorState = std::variant; - // Total number of tables in the composite table - size_t n_tables_; +class AsofJoinNode : public ExecNode { + public: + AsofJoinNode(ExecPlan* plan, NodeVector inputs, std::vector input_labels, + std::vector on_keys, + std::vector> by_keys, + std::vector> input_null_placements, + AsofJoinNodeOptions join_options, std::shared_ptr output_schema, + Ordering output_ordering) + : ExecNode(plan, std::move(inputs), std::move(input_labels), + std::move(output_schema)), + ordering_(std::move(output_ordering)), + on_keys_(std::move(on_keys)), + by_keys_(std::move(by_keys)), + input_null_placements_(std::move(input_null_placements)), + tolerance_(join_options.tolerance, join_options.prefer_earlier_on_tie) {} -#ifndef NDEBUG - // Owning node - AsofJoinNode* node_; -#endif - - static CompositeTable InitUnmaterializedTable( - const std::shared_ptr& schema, - const std::vector>& inputs, arrow::MemoryPool* pool) { - std::unordered_map> dst_to_src; - for (size_t i = 0; i < inputs.size(); i++) { - auto& input = inputs[i]; - for (int src = 0; src < input->get_schema()->num_fields(); src++) { - auto dst = input->MapSrcToDst(src); - if (dst.has_value()) { - dst_to_src[dst.value()] = std::make_pair(static_cast(i), src); - } + Status Init() override { + ARROW_RETURN_NOT_OK(ExecNode::Init()); + input_states_.reserve(inputs_.size()); + rhs_lanes_.reserve(inputs_.size() - 1); + for (size_t i = 0; i < inputs_.size(); ++i) { + input_states_.push_back(std::make_unique( + this, i, inputs_[i], on_keys_[i], by_keys_[i], input_null_placements_[i])); + if (i == 0) { + continue; } - } - return CompositeTable{schema, inputs.size(), dst_to_src, pool}; - } -}; - -// TODO: Currently, AsofJoinNode uses 64-bit hashing which leads to a non-negligible -// probability of collision, which can cause incorrect results when many different by-key -// values are processed. Thus, AsofJoinNode is currently limited to about 100k by-keys for -// guaranteeing this probability is below 1 in a billion. The fix is 128-bit hashing. -// See ARROW-17653 -class AsofJoinNode : public ExecNode { - // 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 UpdateRhs() { - auto& lhs = *state_.at(0); - auto lhs_latest_time = 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]; - - // 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(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; + std::vector payload_columns; + for (int column = 0; column < inputs_[i]->output_schema()->num_fields(); ++column) { + if (column != on_keys_[i] && std::find(by_keys_[i].begin(), by_keys_[i].end(), + column) == by_keys_[i].end()) { + payload_columns.push_back(column); } } + rhs_lanes_.push_back( + std::make_unique(this, i - 1, i, std::move(payload_columns), + tolerance_, plan()->query_context()->memory_pool())); } - return update_state; + return Status::OK(); } - Result> ProcessInner() { - DCHECK(!state_.empty()); - auto& lhs = *state_.at(0); - - // Construct new target table if needed - CompositeTableBuilder dst(state_, output_schema_, - plan()->query_context()->memory_pool(), - DEBUG_ADD(state_.size(), this)); + const char* kind_name() const override { return "AsofJoinNode"; } + const Ordering& ordering() const override { return ordering_; } - // Generate rows into the dst table until we either run out of data or hit the row - // limit, or run out of input - for (;;) { - // If LHS is finished or empty then there's nothing we can do here - if (lhs.Finished() || lhs.Empty()) break; - - 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 (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 (!rhs_update_state.any_advanced) break; // need to wait for new data - } + Status InputReceived(ExecNode* input, ExecBatch batch) override { + if (terminal_.load()) { + return Status::OK(); } - - // Prune memo entries that have expired (to bound memory consumption) - if (!lhs.Empty()) { - for (size_t i = 1; i < state_.size(); ++i) { - OnType ts = tolerance_.Expiry(lhs.GetLatestTime()); - if (ts != TolType::kMinValue) { - state_[i]->RemoveMemoEntriesWithLesserTime(ts); - } - } + if (batch.index == compute::kUnsequencedIndex) { + return Status::Invalid("AsofJoin requires sequenced input"); } - - // Emit the batch - if (dst.empty()) { - return NULLPTR; - } else { - ARROW_ASSIGN_OR_RAISE(auto out, dst.Materialize()); - return out.has_value() ? out.value() : NULLPTR; + auto it = std::find(inputs_.begin(), inputs_.end(), input); + if (it == inputs_.end()) { + return Status::Invalid("AsofJoin received a batch from an unknown input"); } + return input_states_[it - inputs_.begin()]->InsertBatch(std::move(batch)); } -#ifdef ARROW_ENABLE_THREADING - - template - struct Defer { - Callable callable; - explicit Defer(Callable callable) : callable(std::move(callable)) {} - ~Defer() noexcept { callable(); } - }; - - void EndFromProcessThread(Status st = Status::OK()) { - // We must spawn a new task to transfer off the process thread when - // marking this finished. Otherwise there is a chance that doing so could - // mark the plan finished which may destroy the plan which will destroy this - // node which will cause us to join on ourselves. - ARROW_UNUSED( - plan_->query_context()->executor()->Spawn([this, st = std::move(st)]() mutable { - Defer cleanup([this, &st]() { process_task_.MarkFinished(st); }); - if (st.ok()) { - st = output_->InputFinished(this, batches_produced_); - } - for (size_t i = 0; i < state_.size(); ++i) { - const auto& s = state_[i]; - s->ForceShutdown(); - st &= inputs_[i]->StopProducing(); - } - })); - } - - bool CheckEnded() { - if (state_.at(0)->Finished()) { - EndFromProcessThread(); - return false; + Status InputFinished(ExecNode* input, int total_batches) override { + if (terminal_.load()) { + return Status::OK(); } - return true; + if (total_batches < 0) { + return Status::Invalid("AsofJoin input reported a negative batch count"); + } + auto it = std::find(inputs_.begin(), inputs_.end(), input); + if (it == inputs_.end()) { + return Status::Invalid("AsofJoin received completion from an unknown input"); + } + size_t index = static_cast(it - inputs_.begin()); + if (index == 0) { + return LeftInputFinished(total_batches); + } + ARROW_ASSIGN_OR_RAISE(auto task, rhs_lanes_[index - 1]->SetTotal(total_batches)); + if (task) { + return std::move(*task)(); + } + return Status::OK(); } - bool Process() { - std::lock_guard guard(gate_); - if (!CheckEnded()) { - return false; - } + Status StartProducing() override { return Status::OK(); } - // Process batches while we have data - for (;;) { - Result> result = ProcessInner(); - - if (result.ok()) { - auto out_rb = *result; - if (!out_rb) break; - ExecBatch out_b(*out_rb); - out_b.index = batches_produced_++; - DEBUG_SYNC(this, "produce batch ", out_b.index, ":", DEBUG_MANIP(std::endl), - out_rb->ToString(), DEBUG_MANIP(std::endl)); - Status st = output_->InputReceived(this, std::move(out_b)); - if (!st.ok()) { - EndFromProcessThread(std::move(st)); - } - } else { - EndFromProcessThread(result.status()); - return false; - } + void PauseProducing(ExecNode* output, int32_t counter) override { + std::lock_guard lock(coordinator_mutex_); + if (std::holds_alternative(coordinator_) || + counter <= downstream_counter_) { + return; } - - // Report to the output the total batch count, if we've already finished everything - // (there are two places where this can happen: here and InputFinished) - // - // It may happen here in cases where InputFinished was called before we were finished - // producing results (so we didn't know the output size at that time) - if (!CheckEnded()) { - return false; + downstream_counter_ = counter; + if (left_gate_ != LeftGate::Flushing) { + left_gate_ = LeftGate::Paused; } - - // There is no more we can do now but there is still work remaining for later when - // more data arrives. - return true; } - void ProcessThread() { - for (;;) { - if (!process_.WaitAndPop()) { - EndFromProcessThread(); + void ResumeProducing(ExecNode* output, int32_t counter) override { + bool activate = false; + Future<> handoff; + { + std::lock_guard lock(coordinator_mutex_); + if (std::holds_alternative(coordinator_) || + counter <= downstream_counter_) { return; } - if (!Process()) { - return; + downstream_counter_ = counter; + activate = left_gate_ == LeftGate::Paused; + if (left_gate_ != LeftGate::Flushing) { + left_gate_ = LeftGate::Open; + } + if (activate) { + // Resume may be called outside the scheduler by a consuming sink. Reserve a + // scheduler task before releasing the coordinator lock so a concurrent final + // InputFinished cannot end the scheduler before activation is registered. + auto maybe_handoff = + plan()->query_context()->BeginExternalTask("AsofJoinNode::ResumeHandoff"); + if (!maybe_handoff.ok() || !maybe_handoff->is_valid()) { + return; + } + handoff = std::move(*maybe_handoff); } } + if (activate) { + Schedule([this] { return ActivateAfterResume(); }, "AsofJoinNode::Resume"); + handoff.MarkFinished(); + } } - static void ProcessThreadWrapper(AsofJoinNode* node) { node->ProcessThread(); } -#endif - - public: - AsofJoinNode(ExecPlan* plan, NodeVector inputs, std::vector input_labels, - const std::vector& indices_of_on_key, - const std::vector>& indices_of_by_key, - AsofJoinNodeOptions join_options, std::shared_ptr output_schema, - std::vector> key_hashers, bool must_hash, - bool may_rehash); - - Status Init() override { - auto inputs = this->inputs(); - for (size_t i = 0; i < inputs.size(); i++) { - RETURN_NOT_OK(key_hashers_[i]->Init(plan()->query_context()->exec_context(), - inputs[i]->output_schema())); - ARROW_ASSIGN_OR_RAISE( - auto input_state, - InputState::Make(i, tolerance_, must_hash_, may_rehash_, key_hashers_[i].get(), - inputs[i], this, backpressure_counter_, - inputs[i]->output_schema(), indices_of_on_key_[i], - indices_of_by_key_[i])); - state_.push_back(std::move(input_state)); - } - - col_index_t dst_offset = 0; - for (auto& state : state_) - dst_offset = state->InitSrcToDstMapping(dst_offset, !!dst_offset); + Status StopProducingImpl() override { + terminal_.store(true); + { + std::lock_guard lock(coordinator_mutex_); + coordinator_ = Terminal{}; + left_batches_.clear(); + } + for (auto& lane : rhs_lanes_) { + lane->Stop(); + } + return StopInputs(); + } - return Status::OK(); + void Schedule(Task task, std::string_view name = "AsofJoinNode::Lane") { + plan()->query_context()->ScheduleTask(std::move(task), name); } - virtual ~AsofJoinNode() { -#ifdef ARROW_ENABLE_THREADING - PushProcess(false); - if (process_thread_.joinable()) { - process_thread_.join(); + Result> OnSequenced(size_t input_index, + std::shared_ptr batch) { + if (terminal_.load()) { + return Task([this, input_index] { + InputBatchConsumed(input_index); + return Status::OK(); + }); + } + if (input_index == 0) { + ARROW_RETURN_NOT_OK(OnLeftSequenced(std::move(batch))); + return std::nullopt; } -#endif + return rhs_lanes_[input_index - 1]->Enqueue(std::move(batch)); } - const std::vector& indices_of_on_key() { return indices_of_on_key_; } - const std::vector>& indices_of_by_key() { - return indices_of_by_key_; + void InputBatchConsumed(size_t input_index) { + input_states_[input_index]->BatchConsumed().Apply(); } - static Status is_valid_on_field(const std::shared_ptr& field) { - switch (field->type()->id()) { - case Type::INT8: - case Type::INT16: - case Type::INT32: - case Type::INT64: - case Type::UINT8: - case Type::UINT16: - case Type::UINT32: - case Type::UINT64: - case Type::DATE32: - case Type::DATE64: - case Type::TIME32: - case Type::TIME64: - case Type::TIMESTAMP: + bool IsTerminal() const { return terminal_.load(); } + + Status LaneCompleted(size_t lane_index, std::vector values) { + ExecBatch output_batch; + { + std::lock_guard lock(coordinator_mutex_); + auto* matching = std::get_if(&coordinator_); + if (matching == nullptr) { return Status::OK(); - default: - return Status::Invalid("Unsupported type for on-key ", field->name(), " : ", - field->type()->ToString()); + } + if (lane_index >= matching->results.size() || matching->results[lane_index]) { + return Status::Invalid("AsofJoin lane completed an unexpected job"); + } + matching->results[lane_index] = std::move(values); + if (--matching->remaining != 0) { + return Status::OK(); + } + + auto left = matching->left; + std::vector output_values = left->batch.values; + for (auto& lane_result : matching->results) { + DCHECK(lane_result.has_value()); + for (Datum& value : *lane_result) { + output_values.push_back(std::move(value)); + } + } + output_batch = ExecBatch(std::move(output_values), left->batch.length); + output_batch.index = left->batch.index; + coordinator_ = OutputInFlight{std::move(left)}; } + + ARROW_RETURN_NOT_OK(output_->InputReceived(this, std::move(output_batch))); + return OutputDelivered(); } - static Status is_valid_by_field(const std::shared_ptr& field) { + static Status IsValidOnField(const std::shared_ptr& field) { switch (field->type()->id()) { case Type::INT8: case Type::INT16: @@ -1198,18 +709,14 @@ class AsofJoinNode : public ExecNode { case Type::TIME32: case Type::TIME64: case Type::TIMESTAMP: - case Type::STRING: - case Type::LARGE_STRING: - case Type::BINARY: - case Type::LARGE_BINARY: return Status::OK(); default: - return Status::Invalid("Unsupported type for by-key ", field->name(), " : ", + return Status::Invalid("Unsupported type for on-key ", field->name(), " : ", field->type()->ToString()); } } - static Status is_valid_data_field(const std::shared_ptr& field) { + static Status IsValidByField(const std::shared_ptr& field) { switch (field->type()->id()) { case Type::BOOL: case Type::INT8: @@ -1220,8 +727,6 @@ class AsofJoinNode : public ExecNode { case Type::UINT16: case Type::UINT32: case Type::UINT64: - case Type::FLOAT: - case Type::DOUBLE: case Type::DATE32: case Type::DATE64: case Type::TIME32: @@ -1231,90 +736,86 @@ class AsofJoinNode : public ExecNode { case Type::LARGE_STRING: case Type::BINARY: case Type::LARGE_BINARY: + case Type::FIXED_SIZE_BINARY: + case Type::DECIMAL32: + case Type::DECIMAL64: + case Type::DECIMAL128: + case Type::DECIMAL256: return Status::OK(); default: - return Status::Invalid("Unsupported type for data field ", field->name(), " : ", + return Status::Invalid("Unsupported type for by-key ", field->name(), " : ", field->type()->ToString()); } } - /// \brief Make the output schema of an as-of-join node - /// - /// \param[in] input_schema the schema of each input to the node - /// \param[in] indices_of_on_key the on-key index of each input to the node - /// \param[in] indices_of_by_key the by-key indices of each input to the node - static arrow::Result> MakeOutputSchema( - const std::vector> input_schema, - const std::vector& indices_of_on_key, - const std::vector>& indices_of_by_key) { - std::vector> fields; - - size_t n_by = indices_of_by_key.size() == 0 ? 0 : indices_of_by_key[0].size(); - const DataType* on_key_type = NULLPTR; - std::vector by_key_type(n_by, NULLPTR); - // Take all non-key, non-time RHS fields - for (size_t j = 0; j < input_schema.size(); ++j) { - const auto& on_field_ix = indices_of_on_key[j]; - const auto& by_field_ix = indices_of_by_key[j]; - - if ((on_field_ix == -1) || std_has(by_field_ix, -1)) { - return Status::Invalid("Missing join key on table ", j); - } + static Result> MakeOutputSchema( + const std::vector>& input_schemas, + const std::vector& on_keys, + const std::vector>& by_keys) { + if (input_schemas.size() < 2 || input_schemas.size() != on_keys.size() || + input_schemas.size() != by_keys.size()) { + return Status::Invalid("AsofJoin requires matching schemas and keys for at least ", + "two inputs"); + } - const auto& on_field = input_schema[j]->fields()[on_field_ix]; - std::vector by_field(n_by); - for (size_t k = 0; k < n_by; k++) { - by_field[k] = input_schema[j]->fields()[by_field_ix[k]].get(); + const size_t by_key_count = by_keys[0].size(); + const DataType* on_type = nullptr; + std::vector by_types(by_key_count, nullptr); + std::vector> fields; + + for (size_t input_index = 0; input_index < input_schemas.size(); ++input_index) { + const auto& input_schema = input_schemas[input_index]; + col_index_t on_key = on_keys[input_index]; + if (on_key < 0 || on_key >= input_schema->num_fields() || + by_keys[input_index].size() != by_key_count) { + return Status::Invalid("Missing join key on table ", input_index); } - - if (on_key_type == NULLPTR) { - on_key_type = on_field->type().get(); - } else if (*on_key_type != *on_field->type()) { - return Status::Invalid("Expected on-key type ", *on_key_type, " but got ", + const auto& on_field = input_schema->field(on_key); + if (on_type == nullptr) { + ARROW_RETURN_NOT_OK(IsValidOnField(on_field)); + on_type = on_field->type().get(); + } else if (*on_type != *on_field->type()) { + return Status::Invalid("Expected on-key type ", *on_type, " but got ", *on_field->type(), " for field ", on_field->name(), - " in input ", j); + " in input ", input_index); } - for (size_t k = 0; k < n_by; k++) { - if (by_key_type[k] == NULLPTR) { - by_key_type[k] = by_field[k]->type().get(); - } else if (*by_key_type[k] != *by_field[k]->type()) { - return Status::Invalid("Expected by-key type ", *by_key_type[k], " but got ", - *by_field[k]->type(), " for field ", by_field[k]->name(), - " in input ", j); + + for (size_t key_index = 0; key_index < by_key_count; ++key_index) { + col_index_t by_key = by_keys[input_index][key_index]; + if (by_key < 0 || by_key >= input_schema->num_fields()) { + return Status::Invalid("Missing join key on table ", input_index); + } + const auto& by_field = input_schema->field(by_key); + if (by_types[key_index] == nullptr) { + ARROW_RETURN_NOT_OK(IsValidByField(by_field)); + by_types[key_index] = by_field->type().get(); + } else if (*by_types[key_index] != *by_field->type()) { + return Status::Invalid("Expected by-key type ", *by_types[key_index], + " but got ", *by_field->type(), " for field ", + by_field->name(), " in input ", input_index); } } - for (int i = 0; i < input_schema[j]->num_fields(); ++i) { - const auto field = input_schema[j]->field(i); - bool as_output; // true if the field appears as an output - if (i == on_field_ix) { - ARROW_RETURN_NOT_OK(is_valid_on_field(field)); - // Only add on field from the left table - as_output = (j == 0); - } else if (std_has(by_field_ix, i)) { - ARROW_RETURN_NOT_OK(is_valid_by_field(field)); - // Only add by field from the left table - as_output = (j == 0); - } else { - ARROW_RETURN_NOT_OK(is_valid_data_field(field)); - as_output = true; - } - if (as_output) { - fields.push_back(field); + for (int column = 0; column < input_schema->num_fields(); ++column) { + bool is_key = column == on_key || + std::find(by_keys[input_index].begin(), by_keys[input_index].end(), + column) != by_keys[input_index].end(); + if (input_index == 0 || !is_key) { + auto field = input_schema->field(column); + fields.push_back(input_index == 0 ? field : field->WithNullable(true)); } } } - return std::make_shared(fields); + return std::make_shared(std::move(fields)); } - static inline Result FindColIndex(const Schema& schema, - const FieldRef& field_ref, - std::string_view key_kind) { - auto match_res = field_ref.FindOne(schema); - if (!match_res.ok()) { - return Status::Invalid("Bad join key on table : ", match_res.status().message()); + static Result FindColIndex(const Schema& schema, const FieldRef& field_ref, + std::string_view key_kind) { + auto match_result = field_ref.FindOne(schema); + if (!match_result.ok()) { + return Status::Invalid("Bad join key on table : ", match_result.status().message()); } - ARROW_ASSIGN_OR_RAISE(auto match, match_res); + ARROW_ASSIGN_OR_RAISE(auto match, std::move(match_result)); if (match.indices().size() != 1) { return Status::Invalid("AsOfJoinNode does not support a nested ", key_kind, "-key ", field_ref.ToString()); @@ -1324,289 +825,808 @@ class AsofJoinNode : public ExecNode { static Result GetByKeySize( const std::vector& input_keys) { - size_t n_by = 0; - for (size_t i = 0; i < input_keys.size(); ++i) { - const auto& by_key = input_keys[i].by_key; - if (i == 0) { - n_by = by_key.size(); - } else if (n_by != by_key.size()) { + if (input_keys.size() < 2) { + return Status::Invalid("AsofJoin requires at least two inputs"); + } + const size_t size = input_keys[0].by_key.size(); + for (const auto& keys : input_keys) { + if (keys.by_key.size() != size) { return Status::Invalid("inconsistent size of by-key across inputs"); } } - return n_by; + return size; } static Result> GetIndicesOfOnKey( - const std::vector>& input_schema, + const std::vector>& input_schemas, const std::vector& input_keys) { - if (input_schema.size() != input_keys.size()) { + if (input_schemas.size() != input_keys.size()) { return Status::Invalid("mismatching number of input schema and keys"); } - size_t n_input = input_schema.size(); - std::vector indices_of_on_key(n_input); - for (size_t i = 0; i < n_input; ++i) { - const auto& on_key = input_keys[i].on_key; - ARROW_ASSIGN_OR_RAISE(indices_of_on_key[i], - FindColIndex(*input_schema[i], on_key, "on")); + std::vector indices(input_schemas.size()); + for (size_t i = 0; i < input_schemas.size(); ++i) { + ARROW_ASSIGN_OR_RAISE(indices[i], + FindColIndex(*input_schemas[i], input_keys[i].on_key, "on")); } - return indices_of_on_key; + return indices; } static Result>> GetIndicesOfByKey( - const std::vector>& input_schema, + const std::vector>& input_schemas, const std::vector& input_keys) { - if (input_schema.size() != input_keys.size()) { + if (input_schemas.size() != input_keys.size()) { return Status::Invalid("mismatching number of input schema and keys"); } - ARROW_ASSIGN_OR_RAISE(size_t n_by, GetByKeySize(input_keys)); - size_t n_input = input_schema.size(); - std::vector> indices_of_by_key( - n_input, std::vector(n_by)); - for (size_t i = 0; i < n_input; ++i) { - for (size_t k = 0; k < n_by; k++) { - const auto& by_key = input_keys[i].by_key; - ARROW_ASSIGN_OR_RAISE(indices_of_by_key[i][k], - FindColIndex(*input_schema[i], by_key[k], "by")); + ARROW_ASSIGN_OR_RAISE(size_t by_key_count, GetByKeySize(input_keys)); + std::vector> indices(input_schemas.size(), + std::vector(by_key_count)); + for (size_t input_index = 0; input_index < input_schemas.size(); ++input_index) { + for (size_t key_index = 0; key_index < by_key_count; ++key_index) { + ARROW_ASSIGN_OR_RAISE( + indices[input_index][key_index], + FindColIndex(*input_schemas[input_index], + input_keys[input_index].by_key[key_index], "by")); } } - return indices_of_by_key; + return indices; + } + + static Result> ValidateInputOrdering(const ExecNode& input, + col_index_t on_key, + size_t input_index) { + const Ordering& ordering = input.ordering(); + if (ordering.is_unordered()) { + return Status::Invalid("AsofJoin input ", input_index, + " has no meaningful ordering"); + } + if (ordering.is_implicit()) { + return std::nullopt; + } + + DCHECK(!ordering.sort_keys().empty()); + const SortKey& leading_key = ordering.sort_keys().front(); + auto match_result = leading_key.target.FindOne(*input.output_schema()); + if (!match_result.ok()) { + return Status::Invalid( + "AsofJoin input ", input_index, + " has an invalid leading sort key: ", match_result.status().message()); + } + ARROW_ASSIGN_OR_RAISE(auto match, std::move(match_result)); + if (leading_key.order != SortOrder::Ascending || match.indices().size() != 1 || + match.indices()[0] != on_key) { + return Status::Invalid("AsofJoin input ", input_index, + " must be ordered by its ascending on-key"); + } + return ordering.null_placement().value_or(leading_key.null_placement); + } + + static Result NormalizeOutputOrdering(const ExecNode& left) { + const Ordering& ordering = left.ordering(); + if (ordering.is_implicit() || ordering.is_unordered()) { + return ordering; + } + + std::vector sort_keys; + sort_keys.reserve(ordering.sort_keys().size()); + for (const SortKey& sort_key : ordering.sort_keys()) { + ARROW_ASSIGN_OR_RAISE(auto path, sort_key.target.FindOne(*left.output_schema())); + sort_keys.emplace_back(FieldRef(std::move(path)), sort_key.order, + ordering.null_placement().value_or(sort_key.null_placement)); + } + return Ordering(std::move(sort_keys)); } - static arrow::Result Make(ExecPlan* plan, std::vector inputs, - const ExecNodeOptions& options) { - DCHECK_GE(inputs.size(), 2) << "Must have at least two inputs"; + static Result Make(ExecPlan* plan, std::vector inputs, + const ExecNodeOptions& options) { const auto& join_options = checked_cast(options); - ARROW_ASSIGN_OR_RAISE(size_t n_by, GetByKeySize(join_options.input_keys)); - size_t n_input = inputs.size(); - std::vector input_labels(n_input); - std::vector> input_schema(n_input); - for (size_t i = 0; i < n_input; ++i) { + if (inputs.size() < 2 || inputs.size() != join_options.input_keys.size()) { + return Status::Invalid("AsofJoin requires one key specification per input and at ", + "least two inputs"); + } + if (join_options.tolerance.lower > join_options.tolerance.upper) { + return Status::Invalid("AsofJoin tolerance lower bound must not exceed its upper ", + "bound"); + } + ARROW_RETURN_NOT_OK(GetByKeySize(join_options.input_keys).status()); + + std::vector input_labels(inputs.size()); + std::vector> input_schemas(inputs.size()); + for (size_t i = 0; i < inputs.size(); ++i) { input_labels[i] = i == 0 ? "left" : "right_" + ToChars(i); - input_schema[i] = inputs[i]->output_schema(); - } - ARROW_ASSIGN_OR_RAISE(std::vector indices_of_on_key, - GetIndicesOfOnKey(input_schema, join_options.input_keys)); - ARROW_ASSIGN_OR_RAISE(std::vector> indices_of_by_key, - GetIndicesOfByKey(input_schema, join_options.input_keys)); - ARROW_ASSIGN_OR_RAISE( - std::shared_ptr output_schema, - MakeOutputSchema(input_schema, indices_of_on_key, indices_of_by_key)); - - std::vector> key_hashers; - for (size_t i = 0; i < n_input; i++) { - key_hashers.push_back(std::make_unique(i, indices_of_by_key[i])); - } - bool must_hash = - n_by > 1 || - (n_by == 1 && - !is_primitive( - inputs[0]->output_schema()->field(indices_of_by_key[0][0])->type()->id())); - bool may_rehash = n_by == 1 && !must_hash; + input_schemas[i] = inputs[i]->output_schema(); + } + ARROW_ASSIGN_OR_RAISE(auto on_keys, + GetIndicesOfOnKey(input_schemas, join_options.input_keys)); + ARROW_ASSIGN_OR_RAISE(auto by_keys, + GetIndicesOfByKey(input_schemas, join_options.input_keys)); + std::vector> input_null_placements(inputs.size()); + for (size_t i = 0; i < inputs.size(); ++i) { + ARROW_ASSIGN_OR_RAISE(input_null_placements[i], + ValidateInputOrdering(*inputs[i], on_keys[i], i)); + } + ARROW_ASSIGN_OR_RAISE(auto output_ordering, NormalizeOutputOrdering(*inputs.front())); + ARROW_ASSIGN_OR_RAISE(auto output_schema, + MakeOutputSchema(input_schemas, on_keys, by_keys)); return plan->EmplaceNode( - plan, inputs, std::move(input_labels), std::move(indices_of_on_key), - std::move(indices_of_by_key), std::move(join_options), std::move(output_schema), - std::move(key_hashers), must_hash, may_rehash); + plan, std::move(inputs), std::move(input_labels), std::move(on_keys), + std::move(by_keys), std::move(input_null_placements), join_options, + std::move(output_schema), std::move(output_ordering)); } - const char* kind_name() const override { return "AsofJoinNode"; } - const Ordering& ordering() const override { return ordering_; } + private: + void MaybeEnterFlushingUnlocked() { + if (left_total_batches_ && left_received_batches_ == *left_total_batches_) { + left_gate_ = LeftGate::Flushing; + } + } - Status InputReceived(ExecNode* input, ExecBatch batch) override { - // InputReceived may be called after execution was finished. Pushing it to the - // InputState is unnecessary since we're done (and anyway may cause the - // BackPressureController to pause the input, causing a deadlock), so drop it. - if (::arrow::compute::kUnsequencedIndex == batch.index) - return Status::Invalid("AsofJoin requires sequenced input"); + Result> ActivateNextUnlocked() { + DCHECK(std::holds_alternative(coordinator_)); + DCHECK(!left_batches_.empty()); + auto left = std::move(left_batches_.front()); + left_batches_.pop_front(); + coordinator_ = Matching{left, rhs_lanes_.size()}; - if (process_task_.is_finished()) { - DEBUG_SYNC(this, "Input received while done. Short circuiting.", - DEBUG_MANIP(std::endl)); + std::vector tasks; + tasks.reserve(rhs_lanes_.size()); + for (auto& lane : rhs_lanes_) { + ARROW_ASSIGN_OR_RAISE(Task task, lane->Assign(left)); + tasks.push_back(std::move(task)); + } + return tasks; + } + + Status ActivateOrFinishUnlocked(std::vector* tasks, bool* finish) { + if (!std::holds_alternative(coordinator_)) { return Status::OK(); } + if (!left_batches_.empty()) { + if (left_gate_ == LeftGate::Paused) { + return Status::OK(); + } + ARROW_ASSIGN_OR_RAISE(*tasks, ActivateNextUnlocked()); + return Status::OK(); + } + if (left_batches_.empty() && left_total_batches_ && + left_received_batches_ == *left_total_batches_) { + coordinator_ = Terminal{}; + terminal_.store(true); + *finish = true; + } + return Status::OK(); + } - // Get the input - ARROW_DCHECK(std_has(inputs_, input)); - size_t k = std_find(inputs_, input) - inputs_.begin(); + Status ActivateAfterResume() { + std::vector tasks; + bool finish = false; + { + std::lock_guard lock(coordinator_mutex_); + if (std::holds_alternative(coordinator_) || + left_gate_ == LeftGate::Paused) { + return Status::OK(); + } + ARROW_RETURN_NOT_OK(ActivateOrFinishUnlocked(&tasks, &finish)); + } + ScheduleAll(std::move(tasks)); + return finish ? FinishNormally() : Status::OK(); + } - // Put into the sequencing queue - ARROW_RETURN_NOT_OK(state_.at(k)->InsertBatch(std::move(batch))); + Status OnLeftSequenced(std::shared_ptr batch) { + std::vector tasks; + bool finish = false; + bool consume = false; + { + std::lock_guard lock(coordinator_mutex_); + if (std::holds_alternative(coordinator_)) { + consume = true; + } else { + ++left_received_batches_; + if (left_total_batches_ && left_received_batches_ > *left_total_batches_) { + return Status::Invalid( + "AsofJoin left input produced more batches than declared"); + } + left_batches_.push_back(std::move(batch)); + MaybeEnterFlushingUnlocked(); + ARROW_RETURN_NOT_OK(ActivateOrFinishUnlocked(&tasks, &finish)); + } + } + if (consume) { + // Flow-control callbacks may synchronously produce another batch, so keep them + // outside the coordinator mutex. + InputBatchConsumed(0); + return Status::OK(); + } + DCHECK(!finish); + // If this batch opened a generation, post work for every RHS lane. RHS arrival, + // in contrast, can continue inline when it wakes a waiting lane. + ScheduleAll(std::move(tasks)); + return Status::OK(); + } - PushProcess(true); + Status LeftInputFinished(int total_batches) { + std::vector tasks; + bool finish = false; + { + std::lock_guard lock(coordinator_mutex_); + if (std::holds_alternative(coordinator_)) { + return Status::OK(); + } + if (left_total_batches_) { + return *left_total_batches_ == total_batches + ? Status::OK() + : Status::Invalid("AsofJoin left input changed its total batch count"); + } + if (left_received_batches_ > total_batches) { + return Status::Invalid("AsofJoin left input declared too few batches"); + } + left_total_batches_ = total_batches; + MaybeEnterFlushingUnlocked(); + ARROW_RETURN_NOT_OK(ActivateOrFinishUnlocked(&tasks, &finish)); + } + ScheduleAll(std::move(tasks)); + return finish ? FinishNormally() : Status::OK(); + } + + Status OutputDelivered() { + std::shared_ptr left; + std::vector tasks; + bool finish = false; + { + std::lock_guard lock(coordinator_mutex_); + auto* output = std::get_if(&coordinator_); + if (output == nullptr) { + return Status::OK(); + } + left = output->left; + ++batches_produced_; + coordinator_ = WaitingForLeft{}; + ARROW_RETURN_NOT_OK(ActivateOrFinishUnlocked(&tasks, &finish)); + } + InputBatchConsumed(0); + if (finish) { + return FinishNormally(); + } + // An output boundary is the hard yield point between left-batch generations. A + // downstream pause prevents the next one from opening until the LHS is exhausted; + // then Flushing drains the finite tail. + ScheduleAll(std::move(tasks)); return Status::OK(); } - Status InputFinished(ExecNode* input, int total_batches) override { - { - std::lock_guard guard(gate_); - ARROW_DCHECK(std_has(inputs_, input)); - size_t k = std_find(inputs_, input) - inputs_.begin(); - state_.at(k)->set_total_batches(total_batches); + void ScheduleAll(std::vector tasks) { + for (Task& task : tasks) { + Schedule(std::move(task)); + } + } + + Status FinishNormally() { + Status status = output_->InputFinished(this, batches_produced_); + status &= StopInputs(); + return status; + } + + Status StopInputs() { + Status status = Status::OK(); + for (auto& input_state : input_states_) { + input_state->Shutdown(); } - // Trigger a process call - // The reason for this is that there are cases at the end of a table where we don't - // know whether the RHS of the join is up-to-date until we know that the table is - // finished. - PushProcess(true); + for (ExecNode* input : inputs_) { + status &= input->StopProducing(); + } + return status; + } + + const Ordering ordering_; + std::vector on_keys_; + std::vector> by_keys_; + std::vector> input_null_placements_; + Tolerance tolerance_; + std::vector> input_states_; + std::vector> rhs_lanes_; + + std::mutex coordinator_mutex_; + CoordinatorState coordinator_ = WaitingForLeft{}; + LeftGate left_gate_ = LeftGate::Open; + int32_t downstream_counter_ = std::numeric_limits::min(); + std::deque> left_batches_; + int left_received_batches_ = 0; + std::optional left_total_batches_; + int batches_produced_ = 0; + std::atomic terminal_{false}; +}; + +InputState::InputState(AsofJoinNode* node, size_t index, ExecNode* input, + col_index_t on_key, std::vector by_keys, + std::optional null_placement) + : node_(node), + index_(index), + input_(input), + on_key_(on_key), + by_keys_(std::move(by_keys)), + sequencer_(util::SerialSequencingQueue::Make(this)), + null_placement_(null_placement) {} + +Status InputState::InsertBatch(ExecBatch batch) { + if (batch.index == compute::kUnsequencedIndex) { + return Status::Invalid("AsofJoin requires sequenced input"); + } + return sequencer_->InsertBatch(std::move(batch)); +} +Status InputState::Process(ExecBatch batch) { + if (node_->IsTerminal()) { return Status::OK(); } - void PushProcess(bool value) { -#ifdef ARROW_ENABLE_THREADING - process_.Push(value); -#else - if (value) { - ProcessNonThreaded(); - } else if (!process_task_.is_finished()) { - EndFromSingleThread(); - } -#endif - } - -#ifndef ARROW_ENABLE_THREADING - bool ProcessNonThreaded() { - while (!process_task_.is_finished()) { - Result> result = ProcessInner(); - - if (result.ok()) { - auto out_rb = *result; - if (!out_rb) break; - ExecBatch out_b(*out_rb); - out_b.index = batches_produced_++; - DEBUG_SYNC(this, "produce batch ", out_b.index, ":", DEBUG_MANIP(std::endl), - out_rb->ToString(), DEBUG_MANIP(std::endl)); - Status st = output_->InputReceived(this, std::move(out_b)); - if (!st.ok()) { - // this isn't really from a thread, - // but we call through to this for consistency - EndFromSingleThread(std::move(st)); - return false; - } - } else { - // this isn't really from a thread, - // but we call through to this for consistency - EndFromSingleThread(result.status()); - return false; + + ARROW_ASSIGN_OR_RAISE(auto prepared, + PrepareBatch(std::move(batch), on_key_, by_keys_, + node_->plan()->query_context()->exec_context())); + ARROW_RETURN_NOT_OK(ValidateTimes(*prepared)); + + // Only sequenced batches count toward backpressure. Counting physical arrivals can + // deadlock with a reordering input: the high watermark may be reached by later + // batches while that input still holds the missing earlier batch. + FlowAction buffered = BatchBuffered(); + ARROW_ASSIGN_OR_RAISE(auto task, node_->OnSequenced(index_, std::move(prepared))); + buffered.Apply(); + return task ? std::move(*task)() : Status::OK(); +} + +Status InputState::ValidateTimes(const PreparedBatch& batch) { + for (const auto& time : batch.times) { + if (!time) { + if (null_placement_ == NullPlacement::AtStart && last_time_) { + return Status::Invalid("AsofJoin does not allow out-of-order on-key values"); + } + if (null_placement_ == NullPlacement::AtEnd) { + saw_trailing_null_ = true; } + continue; } - auto& lhs = *state_.at(0); - if (lhs.Finished() && !process_task_.is_finished()) { - EndFromSingleThread(Status::OK()); + if ((null_placement_ == NullPlacement::AtEnd && saw_trailing_null_) || + (last_time_ && *time < *last_time_)) { + return Status::Invalid("AsofJoin does not allow out-of-order on-key values"); } + last_time_ = *time; + } + return Status::OK(); +} + +FlowAction InputState::SetUpstreamPausedUnlocked(bool paused) { + if (upstream_paused_ == paused || shutdown_) { + return {}; + } + upstream_paused_ = paused; + return {paused ? FlowAction::Kind::Pause : FlowAction::Kind::Resume, input_, node_, + ++outgoing_counter_}; +} + +FlowAction InputState::BatchBuffered() { + std::lock_guard lock(flow_mutex_); + if (shutdown_) { + return {}; + } + ++buffered_batches_; + return buffered_batches_ >= kHighWatermark ? SetUpstreamPausedUnlocked(true) + : FlowAction{}; +} + +FlowAction InputState::BatchConsumed() { + std::lock_guard lock(flow_mutex_); + if (shutdown_) { + return {}; + } + DCHECK_GT(buffered_batches_, 0); + if (buffered_batches_ > 0) { + --buffered_batches_; + } + return buffered_batches_ <= kLowWatermark ? SetUpstreamPausedUnlocked(false) + : FlowAction{}; +} + +void InputState::Shutdown() { + std::lock_guard lock(flow_mutex_); + if (shutdown_) { + return; + } + shutdown_ = true; + buffered_batches_ = 0; +} + +bool RhsLane::StreamEndedUnlocked() const { + return total_batches_.has_value() && received_batches_ == *total_batches_; +} + +Result> RhsLane::Enqueue(std::shared_ptr batch) { + bool claim = false; + { + std::lock_guard lock(mutex_); + if (phase_ == Phase::Stopped) { + return std::nullopt; + } + if (total_batches_ && received_batches_ >= *total_batches_) { + return Status::Invalid("AsofJoin right input produced more batches than declared"); + } + batches_.push_back(std::move(batch)); + ++received_batches_; + if (phase_ == Phase::Waiting) { + phase_ = Phase::Claimed; + claim = true; + } + } + if (!claim) { + return std::nullopt; + } + return Task([this] { return Run(); }); +} + +Result> RhsLane::SetTotal(int total_batches) { + bool claim = false; + { + std::lock_guard lock(mutex_); + if (phase_ == Phase::Stopped) { + return std::nullopt; + } + if (total_batches_ && *total_batches_ != total_batches) { + return Status::Invalid("AsofJoin right input changed its total batch count"); + } + if (received_batches_ > total_batches) { + return Status::Invalid("AsofJoin right input declared too few batches"); + } + total_batches_ = total_batches; + if (phase_ == Phase::Waiting && StreamEndedUnlocked()) { + phase_ = Phase::Claimed; + claim = true; + } + } + if (!claim) { + return std::nullopt; + } + return Task([this] { return Run(); }); +} + +Result RhsLane::Assign(std::shared_ptr left) { + std::lock_guard lock(mutex_); + if (phase_ == Phase::Stopped) { + return Task([] { return Status::OK(); }); + } + if (phase_ != Phase::NoJob || job_) { + return Status::Invalid("AsofJoin RHS lane was assigned overlapping left batches"); + } + job_ = std::make_shared(std::move(left)); + phase_ = Phase::Claimed; + return Task([this] { return Run(); }); +} + +void RhsLane::Stop() { + std::lock_guard lock(mutex_); + phase_ = Phase::Stopped; +} + +Result RhsLane::PeekNext() { + for (;;) { + if (node_->IsTerminal()) { + return PeekResult{PeekResult::Kind::End, {}}; + } + if (current_batch_) { + if (current_row_ < current_batch_->batch.length) { + return PeekResult{PeekResult::Kind::Row, RowRef{current_batch_, current_row_}}; + } + current_batch_.reset(); + current_row_ = 0; + node_->InputBatchConsumed(input_index_); + } + + { + std::lock_guard lock(mutex_); + if (phase_ == Phase::Stopped) { + return PeekResult{PeekResult::Kind::End, {}}; + } + if (batches_.empty()) { + return PeekResult{ + StreamEndedUnlocked() ? PeekResult::Kind::End : PeekResult::Kind::Blocked, + {}}; + } + current_batch_ = std::move(batches_.front()); + batches_.pop_front(); + } + } +} + +void RhsLane::ConsumeNext() { + DCHECK(current_batch_); + DCHECK_LT(current_row_, current_batch_->batch.length); + if (++current_row_ == current_batch_->batch.length) { + current_batch_.reset(); + current_row_ = 0; + node_->InputBatchConsumed(input_index_); + } +} + +bool RhsLane::WaitOrRetry() { + std::lock_guard lock(mutex_); + if (phase_ == Phase::Stopped) { + return false; + } + if (!batches_.empty() || StreamEndedUnlocked()) { return true; } + phase_ = Phase::Waiting; + return false; +} + +void RhsLane::RememberBackward(const RowRef& row, OnType time) { + const uint64_t hash = row.batch->Hash(row.row); + const uint64_t version = ++next_version_; + auto& candidates = backward_candidates_[hash]; + auto candidate = + std::find_if(candidates.begin(), candidates.end(), + [&](const Candidate& value) { return KeysEqual(row, value.row); }); + if (candidate == candidates.end()) { + candidates.push_back(Candidate{row, version}); + } else { + *candidate = Candidate{row, version}; + } + backward_expiry_.push_back(ExpiryEntry{time, hash, version}); +} - void EndFromSingleThread(Status st = Status::OK()) { - process_task_.MarkFinished(st); - if (st.ok()) { - st = output_->InputFinished(this, batches_produced_); +void RhsLane::ExpireBackward(OnType lower_bound) { + while (!backward_expiry_.empty() && backward_expiry_.front().time < lower_bound) { + ExpiryEntry expired = std::move(backward_expiry_.front()); + backward_expiry_.pop_front(); + auto bucket = backward_candidates_.find(expired.hash); + if (bucket == backward_candidates_.end()) { + continue; } + auto& candidates = bucket->second; + auto candidate = std::find_if( + candidates.begin(), candidates.end(), + [&](const Candidate& value) { return value.version == expired.version; }); + if (candidate != candidates.end()) { + candidates.erase(candidate); + } + if (candidates.empty()) { + backward_candidates_.erase(bucket); + } + } +} - for (size_t i = 0; i < state_.size(); ++i) { - const auto& s = state_[i]; - s->ForceShutdown(); - st &= inputs_[i]->StopProducing(); +void RhsLane::RememberOrdered(const RowRef& row, OnType time) { + const uint64_t hash = row.batch->Hash(row.row); + const uint64_t version = ++next_version_; + auto& bucket = ordered_candidates_[hash]; + auto candidates = + std::find_if(bucket.begin(), bucket.end(), [&](const OrderedCandidates& value) { + return !value.rows.empty() && KeysEqual(row, value.rows.front().row); + }); + if (candidates == bucket.end()) { + bucket.emplace_back(); + candidates = std::prev(bucket.end()); + } + candidates->rows.push_back(OrderedCandidate{row, time, version}); + ordered_expiry_.push_back(ExpiryEntry{time, hash, version}); +} + +void RhsLane::ExpireOrdered(OnType lower_bound) { + while (!ordered_expiry_.empty() && ordered_expiry_.front().time < lower_bound) { + ExpiryEntry expired = std::move(ordered_expiry_.front()); + ordered_expiry_.pop_front(); + auto bucket = ordered_candidates_.find(expired.hash); + if (bucket == ordered_candidates_.end()) { + continue; + } + auto& states = bucket->second; + auto candidates = + std::find_if(states.begin(), states.end(), [&](const OrderedCandidates& value) { + return !value.rows.empty() && value.rows.front().version == expired.version; + }); + if (candidates != states.end()) { + candidates->rows.pop_front(); + if (candidates->rows.empty()) { + states.erase(candidates); + } + } + if (states.empty()) { + ordered_candidates_.erase(bucket); } } +} + +std::optional RhsLane::MatchBackward(const RowRef& key) const { + auto bucket = backward_candidates_.find(key.batch->Hash(key.row)); + if (bucket == backward_candidates_.end()) { + return std::nullopt; + } + auto candidate = + std::find_if(bucket->second.begin(), bucket->second.end(), + [&](const Candidate& value) { return KeysEqual(key, value.row); }); + return candidate == bucket->second.end() ? std::nullopt + : std::optional(candidate->row); +} + +std::optional RhsLane::MatchOrdered(const RowRef& key, OnType left_time) const { + auto bucket = ordered_candidates_.find(key.batch->Hash(key.row)); + if (bucket == ordered_candidates_.end()) { + return std::nullopt; + } + auto candidates = std::find_if( + bucket->second.begin(), bucket->second.end(), [&](const OrderedCandidates& value) { + return !value.rows.empty() && KeysEqual(key, value.rows.front().row); + }); + if (candidates == bucket->second.end()) { + return std::nullopt; + } + + const auto& rows = candidates->rows; + auto later = std::lower_bound(rows.begin(), rows.end(), left_time, + [](const OrderedCandidate& candidate, OnType time) { + return candidate.time < time; + }); + if (later == rows.begin()) { + return later->row; + } + if (later == rows.end()) { + return rows.back().row; + } + if (later->time == left_time) { + return later->row; + } -#endif + const auto earlier = std::prev(later); + const OnType earlier_distance = left_time - earlier->time; + const OnType later_distance = later->time - left_time; + if (earlier_distance < later_distance || + (earlier_distance == later_distance && tolerance_.prefer_earlier_on_tie())) { + return earlier->row; + } + return later->row; +} - Status StartProducing() override { - ARROW_ASSIGN_OR_RAISE(process_task_, plan_->query_context()->BeginExternalTask( - "AsofJoinNode::ProcessThread")); - if (!process_task_.is_valid()) { - // Plan has already aborted. Do not start process thread +Status RhsLane::Run() { + std::shared_ptr job; + { + std::lock_guard lock(mutex_); + if (phase_ != Phase::Claimed || !job_) { return Status::OK(); } -#ifdef ARROW_ENABLE_THREADING - process_thread_ = std::thread(&AsofJoinNode::ProcessThreadWrapper, this); -#endif - return Status::OK(); + job = job_; } - void PauseProducing(ExecNode* output, int32_t counter) override {} - void ResumeProducing(ExecNode* output, int32_t counter) override {} + while (job->left_row < job->left->batch.length) { + if (node_->IsTerminal()) { + return Status::OK(); + } + const int64_t left_row = job->left_row; + const auto& left_time = job->left->times[left_row]; + if (!left_time) { + job->AppendMatch(std::nullopt); + ++job->left_row; + continue; + } - Status StopProducingImpl() override { -#ifdef ARROW_ENABLE_THREADING - process_.Clear(); -#endif - PushProcess(false); - return Status::OK(); + const auto bounds = tolerance_.BoundsFor(*left_time); + if (!bounds) { + job->AppendMatch(std::nullopt); + ++job->left_row; + continue; + } + + const RowRef key{job->left, left_row}; + if (tolerance_.mode() == CandidateMode::Latest) { + ExpireBackward(bounds->lower); + } else { + ExpireOrdered(bounds->lower); + } + + for (;;) { + ARROW_ASSIGN_OR_RAISE(PeekResult next, PeekNext()); + if (next.kind != PeekResult::Kind::Row) { + if (next.kind == PeekResult::Kind::Blocked) { + if (WaitOrRetry()) { + continue; + } + return Status::OK(); + } + break; + } + + const auto& right_time = next.row.batch->times[next.row.row]; + if (!right_time) { + ConsumeNext(); + continue; + } + if (*right_time > bounds->upper) { + break; + } + + if (*right_time >= bounds->lower) { + if (tolerance_.mode() == CandidateMode::Latest) { + RememberBackward(next.row, *right_time); + } else { + RememberOrdered(next.row, *right_time); + } + } + ConsumeNext(); + } + + job->AppendMatch(tolerance_.mode() == CandidateMode::Latest + ? MatchBackward(key) + : MatchOrdered(key, *left_time)); + ++job->left_row; } -#ifndef NDEBUG - std::ostream* GetDebugStream() { return debug_os_; } + ARROW_ASSIGN_OR_RAISE(auto values, Materialize(*job)); + { + std::lock_guard lock(mutex_); + if (phase_ == Phase::Stopped) { + return Status::OK(); + } + if (job_ != job || phase_ != Phase::Claimed) { + return Status::Invalid("AsofJoin RHS lane lost ownership of its job"); + } + job_.reset(); + phase_ = Phase::NoJob; + } + return node_->LaneCompleted(lane_index_, std::move(values)); +} - std::mutex* GetDebugMutex() { return debug_mutex_; } -#endif +Result> RhsLane::Materialize(const Job& job) const { + std::vector values; + values.reserve(payload_columns_.size()); + if (payload_columns_.empty()) { + return values; + } + + const auto& input_schema = node_->inputs()[input_index_]->output_schema(); + for (col_index_t column : payload_columns_) { + ARROW_ASSIGN_OR_RAISE(auto builder, + MakeBuilder(input_schema->field(column)->type(), pool_)); + ARROW_RETURN_NOT_OK(builder->Reserve(job.left->batch.length)); + const ArrayData* current_source = nullptr; + std::optional current_span; + for (const SelectionRun& selection : job.selections) { + if (!selection.source) { + ARROW_RETURN_NOT_OK(builder->AppendNulls(selection.length)); + continue; + } - private: - // Outputs from this node are always in ascending order according to the on key - const Ordering ordering_; - std::vector indices_of_on_key_; - std::vector> indices_of_by_key_; - std::vector> key_hashers_; - bool must_hash_; - bool may_rehash_; - // InputStates - // Each input state corresponds to an input table - std::vector> state_; - std::mutex gate_; - TolType tolerance_; -#ifndef NDEBUG - std::ostream* debug_os_; - std::mutex* debug_mutex_; -#endif - - // Backpressure counter common to all inputs - std::atomic backpressure_counter_; -#ifdef ARROW_ENABLE_THREADING - // Queue for triggering processing of a given input - // (a false value is a poison pill) - ConcurrentQueue process_; - // Worker thread - std::thread process_thread_; -#endif - Future<> process_task_; - - // In-progress batches produced - int batches_produced_ = 0; -}; + const RowRef& match = *selection.source; + const Datum& source = match.batch->batch.values[column]; + if (source.is_scalar()) { + ARROW_RETURN_NOT_OK(builder->AppendScalar(*source.scalar(), selection.length)); + continue; + } + if (!source.is_array()) { + return Status::Invalid( + "AsofJoin RHS payload must be an array or scalar, but got ", + ::arrow::ToString(source.kind())); + } -AsofJoinNode::AsofJoinNode(ExecPlan* plan, NodeVector inputs, - std::vector input_labels, - const std::vector& indices_of_on_key, - const std::vector>& indices_of_by_key, - AsofJoinNodeOptions join_options, - std::shared_ptr output_schema, - std::vector> key_hashers, - bool must_hash, bool may_rehash) - : ExecNode(plan, inputs, input_labels, - /*output_schema=*/std::move(output_schema)), - ordering_({SortKey(indices_of_on_key[0])}), - indices_of_on_key_(std::move(indices_of_on_key)), - indices_of_by_key_(std::move(indices_of_by_key)), - key_hashers_(std::move(key_hashers)), - must_hash_(must_hash), - may_rehash_(may_rehash), - tolerance_(TolType(join_options.tolerance)), -#ifndef NDEBUG - debug_os_(join_options.debug_opts ? join_options.debug_opts->os : nullptr), - debug_mutex_(join_options.debug_opts ? join_options.debug_opts->mutex : nullptr), -#endif - backpressure_counter_(1) -#ifdef ARROW_ENABLE_THREADING - , - process_(), - process_thread_() -#endif -{ - for (auto& key_hasher : key_hashers_) { - key_hasher->node_ = this; + if (source.array().get() != current_source) { + current_source = source.array().get(); + current_span.emplace(*source.array()); + } + Status status = + builder->AppendArraySlice(*current_span, match.row, selection.length); + if (status.IsNotImplemented()) { + auto source_array = MakeArray(source.array()); + for (int64_t row = match.row; row < match.row + selection.length; ++row) { + ARROW_ASSIGN_OR_RAISE(auto scalar, source_array->GetScalar(row)); + ARROW_RETURN_NOT_OK(builder->AppendScalar(*scalar)); + } + } else { + ARROW_RETURN_NOT_OK(status); + } + } + ARROW_ASSIGN_OR_RAISE(auto array, builder->Finish()); + values.emplace_back(std::move(array)); } + return values; } +} // namespace + namespace internal { void RegisterAsofJoinNode(ExecFactoryRegistry* registry) { DCHECK_OK(registry->AddFactory("asofjoin", AsofJoinNode::Make)); @@ -1618,26 +1638,13 @@ namespace asofjoin { Result> MakeOutputSchema( const std::vector>& input_schema, const std::vector& input_keys) { - ARROW_ASSIGN_OR_RAISE(std::vector indices_of_on_key, + ARROW_ASSIGN_OR_RAISE(auto on_keys, AsofJoinNode::GetIndicesOfOnKey(input_schema, input_keys)); - ARROW_ASSIGN_OR_RAISE(std::vector> indices_of_by_key, + ARROW_ASSIGN_OR_RAISE(auto by_keys, AsofJoinNode::GetIndicesOfByKey(input_schema, input_keys)); - return AsofJoinNode::MakeOutputSchema(input_schema, indices_of_on_key, - indices_of_by_key); + return AsofJoinNode::MakeOutputSchema(input_schema, on_keys, by_keys); } } // namespace asofjoin - -#ifndef NDEBUG -std::ostream* GetDebugStream(AsofJoinNode* node) { return node->GetDebugStream(); } - -std::mutex* GetDebugMutex(AsofJoinNode* node) { return node->GetDebugMutex(); } -#endif - -#undef DEBUG_SYNC -#undef DEBUG_MANIP -#undef NDEBUG_EXPLICIT -#undef DEBUG_ADD - } // namespace acero } // namespace arrow diff --git a/cpp/src/arrow/acero/asof_join_node_test.cc b/cpp/src/arrow/acero/asof_join_node_test.cc index 59a9b4ebba12..3a905b552318 100644 --- a/cpp/src/arrow/acero/asof_join_node_test.cc +++ b/cpp/src/arrow/acero/asof_join_node_test.cc @@ -42,13 +42,13 @@ #include "arrow/api.h" #include "arrow/compute/api_scalar.h" #include "arrow/compute/cast.h" +#include "arrow/compute/key_hash_internal.h" #include "arrow/compute/row/row_encoder_internal.h" #include "arrow/compute/test_util_internal.h" #include "arrow/testing/generator.h" #include "arrow/testing/gtest_util.h" #include "arrow/testing/matchers.h" #include "arrow/testing/random.h" -#include "arrow/util/checked_cast.h" #include "arrow/util/logging_internal.h" #include "arrow/util/thread_pool.h" @@ -72,6 +72,7 @@ using compute::Cast; using compute::Divide; using compute::ExecBatchFromJSON; using compute::Multiply; +using compute::SortKey; using compute::Subtract; namespace acero { @@ -254,11 +255,14 @@ void CheckRunOutput(const BatchesWithSchema& l_batches, Declaration join{"asofjoin", join_options}; join.inputs.emplace_back(Declaration{ - "source", SourceNodeOptions{l_batches.schema, l_batches.gen(false, false)}}); + "source", SourceNodeOptions{l_batches.schema, l_batches.gen(false, false), + Ordering::Implicit()}}); join.inputs.emplace_back(Declaration{ - "source", SourceNodeOptions{r0_batches.schema, r0_batches.gen(false, false)}}); + "source", SourceNodeOptions{r0_batches.schema, r0_batches.gen(false, false), + Ordering::Implicit()}}); join.inputs.emplace_back(Declaration{ - "source", SourceNodeOptions{r1_batches.schema, r1_batches.gen(false, false)}}); + "source", SourceNodeOptions{r1_batches.schema, r1_batches.gen(false, false), + Ordering::Implicit()}}); ASSERT_OK_AND_ASSIGN(auto res_table, DeclarationToTable(std::move(join), /*use_threads=*/false)); @@ -278,9 +282,11 @@ void DoInvalidPlanTest(const BatchesWithSchema& l_batches, Declaration join{"asofjoin", join_options}; join.inputs.emplace_back(Declaration{ - "source", SourceNodeOptions{l_batches.schema, l_batches.gen(false, false)}}); + "source", SourceNodeOptions{l_batches.schema, l_batches.gen(false, false), + Ordering::Implicit()}}); join.inputs.emplace_back(Declaration{ - "source", SourceNodeOptions{r_batches.schema, r_batches.gen(false, false)}}); + "source", SourceNodeOptions{r_batches.schema, r_batches.gen(false, false), + Ordering::Implicit()}}); if (fail_on_plan_creation) { EXPECT_RAISES_WITH_MESSAGE_THAT( @@ -1285,6 +1291,10 @@ TRACED_TEST(AsofJoinTest, TestUnsupportedOntype, { field("l_v0", float64())}), schema({field("time", list(int32())), field("key", int32()), field("r0_v0", float32())})); + DoRunInvalidTypeTest( + schema({field("time", boolean()), field("key", int32()), field("l_v0", float64())}), + schema( + {field("time", boolean()), field("key", int32()), field("r0_v0", float32())})); }) TRACED_TEST(AsofJoinTest, TestUnsupportedByType, { @@ -1292,14 +1302,12 @@ TRACED_TEST(AsofJoinTest, TestUnsupportedByType, { field("l_v0", float64())}), schema({field("time", int64()), field("key", list(int32())), field("r0_v0", float32())})); -}) -TRACED_TEST(AsofJoinTest, TestUnsupportedDatatype, { - // List is unsupported - DoRunInvalidTypeTest( - schema({field("time", int64()), field("key", int32()), field("l_v0", float64())}), - schema({field("time", int64()), field("key", int32()), - field("r0_v0", list(int32()))})); + auto dictionary_type = dictionary(int8(), utf8()); + DoRunInvalidTypeTest(schema({field("time", int64()), field("key", dictionary_type), + field("l_v0", float64())}), + schema({field("time", int64()), field("key", dictionary_type), + field("r0_v0", float32())})); }) TRACED_TEST(AsofJoinTest, TestMissingKeys, { @@ -1374,6 +1382,7 @@ TRACED_TEST(AsofJoinTest, TestUnorderedOnKey, { }) struct BackpressureCounters { + std::atomic batch_count = 0; std::atomic pause_count = 0; std::atomic resume_count = 0; }; @@ -1409,7 +1418,10 @@ struct BackpressureCountingNode : public MapNode { } const char* kind_name() const override { return kKindName; } - Result ProcessBatch(ExecBatch batch) override { return batch; } + Result ProcessBatch(ExecBatch batch) override { + ++counters->batch_count; + return batch; + } void PauseProducing(ExecNode* output, int32_t counter) override { ++counters->pause_count; @@ -1487,13 +1499,15 @@ void TestBackpressure(BatchesMaker maker, int batch_size, int num_l_batches, const auto& config = source_configs[i]; if (config.is_delayed) { src_decls.emplace_back( - "source", - SourceNodeOptions(config.schema, MakeDelayedGen(config.batches, "slow_source", - /*delay_sec=*/0.5, - /*noisy=*/false))); + "source", SourceNodeOptions(config.schema, + MakeDelayedGen(config.batches, "slow_source", + /*delay_sec=*/0.5, + /*noisy=*/false), + Ordering::Implicit())); } else { - src_decls.emplace_back("source", - SourceNodeOptions(config.schema, GetGen(config.batches))); + src_decls.emplace_back( + "source", + SourceNodeOptions(config.schema, GetGen(config.batches), Ordering::Implicit())); } bp_options.push_back( std::make_shared(&bp_counters[i])); @@ -1542,7 +1556,7 @@ void TestBackpressure(BatchesMaker maker, int batch_size, int num_l_batches, for (size_t i = 0; i < source_configs.size(); i++) { const auto& counters = bp_counters[i]; if (!source_configs[i].is_gated) { - ASSERT_GE(counters.resume_count, 0); + ASSERT_GT(counters.resume_count, 0); } } } @@ -1570,10 +1584,10 @@ void TestSequencing(BatchesMaker maker, int num_batches, int batch_size) { ASSERT_OK_AND_ASSIGN(auto l_batches, make_shift(l_schema, 0)); ASSERT_OK_AND_ASSIGN(auto r_batches, make_shift(r_schema, 1)); - Declaration l_src = {"source", - SourceNodeOptions(l_schema, l_batches.gen(false, false))}; - Declaration r_src = {"source", - SourceNodeOptions(r_schema, r_batches.gen(false, false))}; + Declaration l_src = {"source", SourceNodeOptions(l_schema, l_batches.gen(false, false), + Ordering::Implicit())}; + Declaration r_src = {"source", SourceNodeOptions(r_schema, r_batches.gen(false, false), + Ordering::Implicit())}; Declaration asofjoin = { "asofjoin", {l_src, r_src}, GetRepeatedOptions(2, "time", {"key"}, 1000)}; @@ -1611,10 +1625,10 @@ void TestSchemaResolution(BatchesMaker maker, int num_batches, int batch_size) { ASSERT_OK_AND_ASSIGN(auto l_batches, make_shift(l_schema, 0)); ASSERT_OK_AND_ASSIGN(auto r_batches, make_shift(r_schema, 1)); - Declaration l_src = {"source", - SourceNodeOptions(l_schema, l_batches.gen(false, false))}; - Declaration r_src = {"source", - SourceNodeOptions(r_schema, r_batches.gen(false, false))}; + Declaration l_src = {"source", SourceNodeOptions(l_schema, l_batches.gen(false, false), + Ordering::Implicit())}; + Declaration r_src = {"source", SourceNodeOptions(r_schema, r_batches.gen(false, false), + Ordering::Implicit())}; Declaration l_project = { "project", {std::move(l_src)}, @@ -1677,7 +1691,6 @@ T GetEnvValue(const std::string& var, T default_value) { } // namespace TEST(AsofJoinTest, BackpressureWithBatchesGen) { - GTEST_SKIP() << "Skipping - see GH-36331"; int num_batches = GetEnvValue("ARROW_BACKPRESSURE_DEMO_NUM_BATCHES", 20); int batch_size = GetEnvValue("ARROW_BACKPRESSURE_DEMO_BATCH_SIZE", 1); return TestBackpressure(MakeIntegerBatchGenForTest, /*batch_size=*/batch_size, @@ -1686,6 +1699,853 @@ TEST(AsofJoinTest, BackpressureWithBatchesGen) { /*slow_r0=*/false); } +namespace { + +class PausingSinkConsumer : public SinkNodeConsumer { + public: + explicit PausingSinkConsumer(bool pause_after_first = true) + : pause_after_first_(pause_after_first) {} + + Status Init(const std::shared_ptr&, BackpressureControl* control, + ExecPlan*) override { + control_.store(control); + return Status::OK(); + } + + Status Consume(ExecBatch) override { + if (batches_received_.fetch_add(1) == 0 && pause_after_first_) { + Pause(); + } + return Status::OK(); + } + + Future<> Finish() override { + finished_.MarkFinished(); + return finished_; + } + + void Pause() { + auto* control = control_.load(); + DCHECK_NE(control, nullptr); + control->Pause(); + paused_.store(true); + } + + void Resume() { + auto* control = control_.load(); + DCHECK_NE(control, nullptr); + control->Resume(); + paused_.store(false); + } + + int batches_received() const { return batches_received_.load(); } + bool is_initialized() const { return control_.load() != nullptr; } + bool is_paused() const { return paused_.load(); } + const Future<>& finished() const { return finished_; } + + private: + bool pause_after_first_; + std::atomic batches_received_{0}; + std::atomic paused_{false}; + std::atomic control_{nullptr}; + Future<> finished_ = Future<>::Make(); +}; + +} // namespace + +TEST(AsofJoinTest, LeftBatchIsOutputBackbone) { + auto left_schema = schema({field("on", int64(), /*nullable=*/false), + field("label", utf8(), /*nullable=*/false)}); + auto right_schema = schema({field("on", int64(), /*nullable=*/false), + field("value", int32(), /*nullable=*/false)}); + + ExecBatch first( + {ArrayFromJSON(int64(), "[1, 2]"), std::make_shared("first")}, 2); + ExecBatch empty({ArrayFromJSON(int64(), "[]"), std::make_shared("empty")}, + 0); + ExecBatch last({ArrayFromJSON(int64(), "[3]"), std::make_shared("last")}, + 1); + ExecBatch right({ArrayFromJSON(int64(), "[1, 2, 3]"), std::make_shared(7)}, + 3); + + for (bool use_threads : {false, true}) { + SCOPED_TRACE(use_threads ? "threaded" : "serial"); + Declaration left_source{"exec_batch_source", ExecBatchSourceNodeOptions( + left_schema, {first, empty, last})}; + Declaration right_source{"exec_batch_source", + ExecBatchSourceNodeOptions(right_schema, {right})}; + QueryOptions query_options; + query_options.use_threads = use_threads; + query_options.sequence_output = false; + Declaration join{"asofjoin", + {std::move(left_source), std::move(right_source)}, + GetRepeatedOptions(2, "on", {}, 0)}; + ASSERT_OK_AND_ASSIGN(auto result, + DeclarationToExecBatches(std::move(join), query_options)); + + ASSERT_EQ(result.batches.size(), 3U); + ASSERT_FALSE(result.schema->field(0)->nullable()); + ASSERT_FALSE(result.schema->field(1)->nullable()); + ASSERT_TRUE(result.schema->field(2)->nullable()); + + const std::vector expected_lengths = {2, 0, 1}; + const std::vector expected_labels = {"first", "empty", "last"}; + for (size_t i = 0; i < result.batches.size(); ++i) { + const ExecBatch& batch = result.batches[i]; + EXPECT_EQ(batch.index, static_cast(i)); + EXPECT_EQ(batch.length, expected_lengths[i]); + ASSERT_TRUE(batch.values[1].is_scalar()); + EXPECT_EQ(batch.values[1].scalar_as().view(), expected_labels[i]); + ASSERT_TRUE(batch.values[2].is_array()); + } + AssertArraysEqual(*ArrayFromJSON(int32(), "[7, 7]"), + *result.batches[0].values[2].make_array()); + AssertArraysEqual(*ArrayFromJSON(int32(), "[]"), + *result.batches[1].values[2].make_array()); + AssertArraysEqual(*ArrayFromJSON(int32(), "[7]"), + *result.batches[2].values[2].make_array()); + } +} + +TEST(AsofJoinTest, ScalarColumnsAreFirstClass) { + auto left_schema = + schema({field("on", int64()), field("key", utf8()), field("left_value", int32())}); + auto right_schema = + schema({field("on", int64()), field("key", utf8()), field("right_value", int32())}); + ExecBatch left({std::make_shared(5), std::make_shared("a"), + std::make_shared(7)}, + 3); + ExecBatch right({std::make_shared(5), std::make_shared("a"), + std::make_shared(9)}, + 2); + + for (bool use_threads : {false, true}) { + SCOPED_TRACE(use_threads ? "threaded" : "serial"); + Declaration left_source{"exec_batch_source", + ExecBatchSourceNodeOptions(left_schema, {left})}; + Declaration right_source{"exec_batch_source", + ExecBatchSourceNodeOptions(right_schema, {right})}; + AsofJoinNodeOptions options({{{"on"}, {{"key"}}}, {{"on"}, {{"key"}}}}, 0); + Declaration join{"asofjoin", + {std::move(left_source), std::move(right_source)}, + std::move(options)}; + ASSERT_OK_AND_ASSIGN(auto result, + DeclarationToExecBatches(std::move(join), use_threads)); + + ASSERT_EQ(result.batches.size(), 1U); + const ExecBatch& output = result.batches[0]; + ASSERT_EQ(output.length, 3); + for (size_t column = 0; column < left.values.size(); ++column) { + ASSERT_TRUE(output.values[column].is_scalar()); + EXPECT_EQ(output.values[column].scalar().get(), left.values[column].scalar().get()); + } + ASSERT_TRUE(output.values[3].is_array()); + AssertArraysEqual(*ArrayFromJSON(int32(), "[9, 9, 9]"), + *output.values[3].make_array()); + } +} + +TEST(AsofJoinTest, MixedScalarAndArrayByKeys) { + auto left_schema = schema({field("on", int64()), field("scalar_key", utf8()), + field("array_key", int32()), field("left_value", int32())}); + auto right_schema = + schema({field("on", int64()), field("scalar_key", utf8()), + field("array_key", int32()), field("right_value", int32())}); + ExecBatch left( + {ArrayFromJSON(int64(), "[1, 2, 3]"), std::make_shared("a"), + ArrayFromJSON(int32(), "[1, 2, 3]"), std::make_shared(7)}, + 3); + ExecBatch right( + {ArrayFromJSON(int64(), "[1, 2, 3]"), std::make_shared("a"), + ArrayFromJSON(int32(), "[1, 99, 3]"), ArrayFromJSON(int32(), "[10, 20, 30]")}, + 3); + + for (bool use_threads : {false, true}) { + SCOPED_TRACE(use_threads ? "threaded" : "serial"); + Declaration left_source{"exec_batch_source", + ExecBatchSourceNodeOptions(left_schema, {left})}; + Declaration right_source{"exec_batch_source", + ExecBatchSourceNodeOptions(right_schema, {right})}; + AsofJoinNodeOptions options({{{"on"}, {{"scalar_key"}, {"array_key"}}}, + {{"on"}, {{"scalar_key"}, {"array_key"}}}}, + 0); + Declaration join{"asofjoin", + {std::move(left_source), std::move(right_source)}, + std::move(options)}; + ASSERT_OK_AND_ASSIGN(auto result, + DeclarationToExecBatches(std::move(join), use_threads)); + + ASSERT_EQ(result.batches.size(), 1U); + ASSERT_TRUE(result.batches[0].values[4].is_array()); + AssertArraysEqual(*ArrayFromJSON(int32(), "[10, null, 30]"), + *result.batches[0].values[4].make_array()); + } +} + +TEST(AsofJoinTest, HashCollisionsUseExactByKeys) { + auto left_schema = + schema({field("on", int64()), field("key_a", int64()), field("key_b", int64())}); + auto right_schema = schema({field("on", int64()), field("key_a", int64()), + field("key_b", int64()), field("value", int32())}); + ExecBatch left = ExecBatchFromJSON({int64(), int64(), int64()}, + "[[1, 0, 0], [1, 1, 383703870779339957]]"); + ExecBatch right = + ExecBatchFromJSON({int64(), int64(), int64(), int32()}, + "[[1, 0, 0, 100], [1, 1, 383703870779339957, 200]]"); + ExecBatch expected = + ExecBatchFromJSON({int64(), int64(), int64(), int32()}, + "[[1, 0, 0, 100], [1, 1, 383703870779339957, 200]]"); + + ExecBatch colliding_keys({left.values[1], left.values[2]}, left.length); + std::vector hashes(left.length); + std::vector column_arrays; + util::TempVectorStack temp_stack; + ASSERT_OK(temp_stack.Init(default_memory_pool(), + compute::Hashing64::kHashBatchTempStackUsage)); + ASSERT_OK(compute::Hashing64::HashBatch( + colliding_keys, hashes.data(), column_arrays, + internal::CpuInfo::GetInstance()->hardware_flags(), &temp_stack, + /*start_row=*/0, left.length)); + ASSERT_EQ(hashes[0], hashes[1]); + + for (int64_t tolerance : {int64_t{0}, int64_t{1}}) { + for (bool use_threads : {false, true}) { + SCOPED_TRACE(std::string(tolerance == 0 ? "backward" : "forward") + + (use_threads ? " threaded" : " serial")); + Declaration left_source{"exec_batch_source", + ExecBatchSourceNodeOptions(left_schema, {left})}; + Declaration right_source{"exec_batch_source", + ExecBatchSourceNodeOptions(right_schema, {right})}; + AsofJoinNodeOptions options( + {{{"on"}, {{"key_a"}, {"key_b"}}}, {{"on"}, {{"key_a"}, {"key_b"}}}}, + tolerance); + Declaration join{"asofjoin", + {std::move(left_source), std::move(right_source)}, + std::move(options)}; + ASSERT_OK_AND_ASSIGN(auto result, + DeclarationToExecBatches(std::move(join), use_threads)); + ASSERT_EQ(result.batches.size(), 1U); + AssertExecBatchesEqual(result.schema, {expected}, result.batches); + } + } +} + +TEST(AsofJoinTest, SupportsAdditionalFlatByKeyTypes) { + struct TestCase { + std::shared_ptr type; + std::string_view left_keys; + std::string_view right_keys; + }; + const std::vector test_cases = { + {boolean(), "[false, true, false, true]", "[false, false, true, true]"}, + {fixed_size_binary(3), R"(["aaa", "bbb", "aaa", "bbb"])", + R"(["aaa", "aaa", "bbb", "bbb"])"}, + {decimal32(9, 2), R"(["1.00", "2.00", "1.00", "2.00"])", + R"(["1.00", "1.00", "2.00", "2.00"])"}, + {decimal64(18, 2), R"(["1.00", "2.00", "1.00", "2.00"])", + R"(["1.00", "1.00", "2.00", "2.00"])"}, + {decimal128(20, 2), R"(["1.00", "2.00", "1.00", "2.00"])", + R"(["1.00", "1.00", "2.00", "2.00"])"}, + {decimal256(40, 2), R"(["1.00", "2.00", "1.00", "2.00"])", + R"(["1.00", "1.00", "2.00", "2.00"])"}, + }; + + for (const TestCase& test_case : test_cases) { + SCOPED_TRACE(test_case.type->ToString()); + auto left_schema = schema({field("on", int64()), field("key", test_case.type), + field("left_value", int32())}); + auto right_schema = schema({field("on", int64()), field("key", test_case.type), + field("right_value", int32())}); + ExecBatch left({ArrayFromJSON(int64(), "[1, 2, 3, 4]"), + ArrayFromJSON(test_case.type, test_case.left_keys), + ArrayFromJSON(int32(), "[10, 20, 30, 40]")}, + 4); + ExecBatch right({ArrayFromJSON(int64(), "[1, 2, 3, 4]"), + ArrayFromJSON(test_case.type, test_case.right_keys), + ArrayFromJSON(int32(), "[100, 200, 300, 400]")}, + 4); + + for (int64_t tolerance : {int64_t{-10}, int64_t{10}}) { + SCOPED_TRACE(tolerance < 0 ? "backward" : "forward"); + const char* expected_right = + tolerance < 0 ? "[100, null, 200, 400]" : "[100, 300, null, 400]"; + ExecBatch expected({left.values[0], left.values[1], left.values[2], + ArrayFromJSON(int32(), expected_right)}, + left.length); + + for (bool use_threads : {false, true}) { + SCOPED_TRACE(use_threads ? "threaded" : "serial"); + Declaration left_source{"exec_batch_source", + ExecBatchSourceNodeOptions(left_schema, {left})}; + Declaration right_source{"exec_batch_source", + ExecBatchSourceNodeOptions(right_schema, {right})}; + AsofJoinNodeOptions options({{{"on"}, {{"key"}}}, {{"on"}, {{"key"}}}}, + tolerance); + Declaration join{"asofjoin", + {std::move(left_source), std::move(right_source)}, + std::move(options)}; + ASSERT_OK_AND_ASSIGN(auto result, + DeclarationToExecBatches(std::move(join), use_threads)); + ASSERT_EQ(result.batches.size(), 1U); + AssertExecBatchesEqual(result.schema, {expected}, result.batches); + } + } + } +} + +TEST(AsofJoinTest, MaterializesGenericRhsPayloads) { + auto dictionary_type = dictionary(int8(), utf8()); + ASSERT_OK_AND_ASSIGN( + auto first_dictionary, + DictionaryArray::FromArrays(dictionary_type, ArrayFromJSON(int8(), "[0]"), + ArrayFromJSON(utf8(), R"(["a"])"))); + ASSERT_OK_AND_ASSIGN( + auto second_dictionary, + DictionaryArray::FromArrays(dictionary_type, ArrayFromJSON(int8(), "[0, 1]"), + ArrayFromJSON(utf8(), R"(["b", "a"])"))); + auto fixed_list_type = fixed_size_list(int32(), 2); + + auto left_schema = schema({field("on", int64())}); + auto right_schema = + schema({field("on", int64()), field("list_value", list(int32()), false), + field("fixed_value", fixed_list_type, false), + field("dictionary_value", dictionary_type, false)}); + ExecBatch left = ExecBatchFromJSON({int64()}, "[[1], [2], [4]]"); + ExecBatch first_right( + {ArrayFromJSON(int64(), "[1]"), ArrayFromJSON(list(int32()), "[[1, 2]]"), + ArrayFromJSON(fixed_list_type, "[[10, 11]]"), first_dictionary}, + 1); + ExecBatch second_right( + {ArrayFromJSON(int64(), "[2, 3]"), ArrayFromJSON(list(int32()), "[[3], [4, 5]]"), + ArrayFromJSON(fixed_list_type, "[[20, 21], [30, 31]]"), second_dictionary}, + 2); + + for (bool use_threads : {false, true}) { + SCOPED_TRACE(use_threads ? "threaded" : "serial"); + Declaration left_source{"exec_batch_source", + ExecBatchSourceNodeOptions(left_schema, {left})}; + Declaration right_source{ + "exec_batch_source", + ExecBatchSourceNodeOptions(right_schema, {first_right, second_right})}; + Declaration join{"asofjoin", + {std::move(left_source), std::move(right_source)}, + GetRepeatedOptions(2, "on", {}, 0)}; + ASSERT_OK_AND_ASSIGN(auto result, + DeclarationToExecBatches(std::move(join), use_threads)); + ASSERT_EQ(result.batches.size(), 1U); + const ExecBatch& output = result.batches[0]; + AssertArraysEqual(*ArrayFromJSON(list(int32()), "[[1, 2], [3], null]"), + *output.values[1].make_array()); + AssertArraysEqual(*ArrayFromJSON(fixed_list_type, "[[10, 11], [20, 21], null]"), + *output.values[2].make_array()); + ASSERT_OK_AND_ASSIGN(Datum decoded_dictionary, Cast(output.values[3], utf8())); + AssertArraysEqual(*ArrayFromJSON(utf8(), R"(["a", "b", null])"), + *decoded_dictionary.make_array()); + } +} + +TEST(AsofJoinTest, SignedTimesNullsAndExtremeTolerance) { + auto left_schema = schema({field("on", int64())}); + auto right_schema = schema({field("on", int64()), field("value", int32())}); + + ExecBatch left = ExecBatchFromJSON({int64()}, "[[-2], [-1], [0], [1], [null]]"); + ExecBatch right = ExecBatchFromJSON({int64(), int32()}, + "[[-3, 300], [-1, 100], [0, 0], [null, 999]]"); + ExecBatch expected = ExecBatchFromJSON( + {int64(), int32()}, "[[-2, 300], [-1, 100], [0, 0], [1, 0], [null, null]]"); + + for (bool use_threads : {false, true}) { + SCOPED_TRACE(use_threads ? "threaded" : "serial"); + Declaration left_source{"exec_batch_source", + ExecBatchSourceNodeOptions(left_schema, {left})}; + Declaration right_source{"exec_batch_source", + ExecBatchSourceNodeOptions(right_schema, {right})}; + AsofJoinNodeOptions options({{{"on"}, {}}, {{"on"}, {}}}, -1); + Declaration join{"asofjoin", {left_source, right_source}, options}; + ASSERT_OK_AND_ASSIGN(auto result, + DeclarationToExecBatches(std::move(join), use_threads)); + AssertExecBatchesEqualIgnoringOrder(result.schema, {expected}, result.batches); + + ExecBatch extreme_left = ExecBatchFromJSON({int64()}, "[[0]]"); + ExecBatch extreme_right = + ExecBatchFromJSON({int64(), int32()}, "[[-9223372036854775808, 9]]"); + ExecBatch extreme_expected = ExecBatchFromJSON({int64(), int32()}, "[[0, 9]]"); + Declaration extreme_left_source{ + "exec_batch_source", ExecBatchSourceNodeOptions(left_schema, {extreme_left})}; + Declaration extreme_right_source{ + "exec_batch_source", ExecBatchSourceNodeOptions(right_schema, {extreme_right})}; + AsofJoinNodeOptions extreme_options({{{"on"}, {}}, {{"on"}, {}}}, + std::numeric_limits::min()); + Declaration extreme_join{ + "asofjoin", {extreme_left_source, extreme_right_source}, extreme_options}; + ASSERT_OK_AND_ASSIGN(auto extreme_result, + DeclarationToExecBatches(std::move(extreme_join), use_threads)); + AssertExecBatchesEqualIgnoringOrder(extreme_result.schema, {extreme_expected}, + extreme_result.batches); + } +} + +TEST(AsofJoinTest, ToleranceRangeExcludesExactMatches) { + auto left_schema = schema({field("on", int64()), field("key", int64())}); + auto right_schema = + schema({field("on", int64()), field("key", int64()), field("value", int32())}); + ExecBatch left = + ExecBatchFromJSON({int64(), int64()}, "[[10, 1], [20, 1], [30, 1], [40, 1]]"); + ExecBatch right = ExecBatchFromJSON({int64(), int64(), int32()}, + "[[9, 1, 1], [12, 1, 2], [30, 1, 3], [41, 1, 4]]"); + + struct TestCase { + AsofJoinNodeOptions::ToleranceRange tolerance; + std::string_view expected; + }; + const std::vector test_cases = { + {{-10, -1}, "[[10, 1, 1], [20, 1, 2], [30, 1, null], [40, 1, 3]]"}, + {{1, 10}, "[[10, 1, 2], [20, 1, 3], [30, 1, null], [40, 1, 4]]"}, + }; + + for (const auto& test_case : test_cases) { + for (bool use_threads : {false, true}) { + SCOPED_TRACE(use_threads ? "threaded" : "serial"); + Declaration left_source{"exec_batch_source", + ExecBatchSourceNodeOptions(left_schema, {left})}; + Declaration right_source{"exec_batch_source", + ExecBatchSourceNodeOptions(right_schema, {right})}; + AsofJoinNodeOptions options({{{"on"}, {{"key"}}}, {{"on"}, {{"key"}}}}, + test_case.tolerance); + Declaration join{"asofjoin", {left_source, right_source}, options}; + ASSERT_OK_AND_ASSIGN(auto result, + DeclarationToExecBatches(std::move(join), use_threads)); + ExecBatch expected = + ExecBatchFromJSON({int64(), int64(), int32()}, test_case.expected); + AssertExecBatchesEqual(result.schema, {expected}, result.batches); + } + } +} + +TEST(AsofJoinTest, ToleranceRangeSelectsNearestAndBreaksTies) { + auto left_schema = schema({field("on", int64()), field("key", int64())}); + auto right_schema = + schema({field("on", int64()), field("key", int64()), field("value", int32())}); + ExecBatch left = + ExecBatchFromJSON({int64(), int64()}, "[[10, 1], [20, 1], [30, 1], [40, 1]]"); + ExecBatch right = + ExecBatchFromJSON({int64(), int64(), int32()}, + "[[8, 1, 8], [12, 1, 12], [18, 1, 18], [22, 1, 22], [30, 1, 30], " + "[38, 1, 38], [41, 1, 41]]"); + + for (bool prefer_earlier_on_tie : {true, false}) { + for (bool use_threads : {false, true}) { + SCOPED_TRACE(prefer_earlier_on_tie ? "prefer earlier" : "prefer later"); + SCOPED_TRACE(use_threads ? "threaded" : "serial"); + Declaration left_source{"exec_batch_source", + ExecBatchSourceNodeOptions(left_schema, {left})}; + Declaration right_source{"exec_batch_source", + ExecBatchSourceNodeOptions(right_schema, {right})}; + AsofJoinNodeOptions options({{{"on"}, {{"key"}}}, {{"on"}, {{"key"}}}}, {-3, 3}, + prefer_earlier_on_tie); + Declaration join{"asofjoin", {left_source, right_source}, options}; + ASSERT_OK_AND_ASSIGN(auto result, + DeclarationToExecBatches(std::move(join), use_threads)); + const std::string_view expected_json = + prefer_earlier_on_tie ? "[[10, 1, 8], [20, 1, 18], [30, 1, 30], [40, 1, 41]]" + : "[[10, 1, 12], [20, 1, 22], [30, 1, 30], [40, 1, 41]]"; + ExecBatch expected = ExecBatchFromJSON({int64(), int64(), int32()}, expected_json); + AssertExecBatchesEqual(result.schema, {expected}, result.batches); + } + } +} + +TEST(AsofJoinTest, RejectsInvalidToleranceRange) { + auto left_schema = schema({field("on", int64()), field("key", int64())}); + auto right_schema = + schema({field("on", int64()), field("key", int64()), field("value", int32())}); + AsofJoinNodeOptions options({{{"on"}, {{"key"}}}, {{"on"}, {{"key"}}}}, {1, -1}); + DoRunInvalidPlanTest(left_schema, right_schema, options, + "tolerance lower bound must not exceed its upper bound"); +} + +TEST(AsofJoinTest, ToleranceRangeHandlesOnKeyLimits) { + auto left_schema = schema({field("on", int64())}); + auto right_schema = schema({field("on", int64()), field("value", int32())}); + + struct TestCase { + std::string_view left; + std::string_view right; + AsofJoinNodeOptions::ToleranceRange tolerance; + std::string_view expected; + }; + const std::vector test_cases = { + {"[[-9223372036854775808], [-9223372036854775807]]", + "[[-9223372036854775808, 1]]", + {-1, -1}, + "[[-9223372036854775808, null], [-9223372036854775807, 1]]"}, + {"[[9223372036854775806], [9223372036854775807]]", + "[[9223372036854775807, 1]]", + {1, 1}, + "[[9223372036854775806, 1], [9223372036854775807, null]]"}, + }; + + for (const auto& test_case : test_cases) { + for (bool use_threads : {false, true}) { + SCOPED_TRACE(use_threads ? "threaded" : "serial"); + ExecBatch left = ExecBatchFromJSON({int64()}, test_case.left); + ExecBatch right = ExecBatchFromJSON({int64(), int32()}, test_case.right); + Declaration left_source{"exec_batch_source", + ExecBatchSourceNodeOptions(left_schema, {left})}; + Declaration right_source{"exec_batch_source", + ExecBatchSourceNodeOptions(right_schema, {right})}; + AsofJoinNodeOptions options({{{"on"}, {}}, {{"on"}, {}}}, test_case.tolerance); + Declaration join{"asofjoin", {left_source, right_source}, options}; + ASSERT_OK_AND_ASSIGN(auto result, + DeclarationToExecBatches(std::move(join), use_threads)); + ExecBatch expected = ExecBatchFromJSON({int64(), int32()}, test_case.expected); + AssertExecBatchesEqual(result.schema, {expected}, result.batches); + } + } +} + +TEST(AsofJoinTest, ImplicitOrderingIgnoresNullPlacement) { + auto left_schema = schema({field("on", int64())}); + auto right_schema = schema({field("on", int64()), field("value", int32())}); + ExecBatch left = ExecBatchFromJSON({int64()}, "[[null], [0], [null], [1], [null]]"); + ExecBatch right = ExecBatchFromJSON( + {int64(), int32()}, "[[null, 9], [0, 10], [null, 9], [1, 11], [null, 9]]"); + ExecBatch expected = ExecBatchFromJSON( + {int64(), int32()}, "[[null, null], [0, 10], [null, null], [1, 11], [null, null]]"); + + for (bool use_threads : {false, true}) { + SCOPED_TRACE(use_threads ? "threaded" : "serial"); + Declaration left_source{"exec_batch_source", + ExecBatchSourceNodeOptions(left_schema, {left})}; + Declaration right_source{"exec_batch_source", + ExecBatchSourceNodeOptions(right_schema, {right})}; + Declaration join{"asofjoin", + {std::move(left_source), std::move(right_source)}, + GetRepeatedOptions(2, "on", {}, 0)}; + ASSERT_OK_AND_ASSIGN(auto result, + DeclarationToExecBatches(std::move(join), use_threads)); + AssertExecBatchesEqualIgnoringOrder(result.schema, {expected}, result.batches); + } +} + +TEST(AsofJoinTest, ExplicitOrderingControlsNullPlacement) { + auto left_schema = schema({field("on", int64())}); + auto right_schema = schema({field("on", int64()), field("value", int32())}); + auto make_source = [](std::shared_ptr schema, std::vector batches, + compute::NullPlacement null_placement) { + BatchesWithSchema input{std::move(batches), schema}; + return Declaration{ + "source", + SourceNodeOptions( + std::move(schema), input.gen(/*parallel=*/false, /*slow=*/false), + Ordering({SortKey("on", compute::SortOrder::Ascending, null_placement)}))}; + }; + + ExecBatch left = ExecBatchFromJSON({int64()}, "[[null], [0], [1]]"); + ExecBatch right = + ExecBatchFromJSON({int64(), int32()}, "[[0, 10], [1, 11], [null, 9]]"); + ExecBatch expected = + ExecBatchFromJSON({int64(), int32()}, "[[null, null], [0, 10], [1, 11]]"); + for (bool use_threads : {false, true}) { + SCOPED_TRACE(use_threads ? "threaded" : "serial"); + Declaration join{"asofjoin", + {make_source(left_schema, {left}, compute::NullPlacement::AtStart), + make_source(right_schema, {right}, compute::NullPlacement::AtEnd)}, + GetRepeatedOptions(2, "on", {}, 0)}; + ASSERT_OK_AND_ASSIGN(auto result, + DeclarationToExecBatches(std::move(join), use_threads)); + AssertExecBatchesEqualIgnoringOrder(result.schema, {expected}, result.batches); + } + + for (auto null_placement : + {compute::NullPlacement::AtStart, compute::NullPlacement::AtEnd}) { + SCOPED_TRACE(null_placement == compute::NullPlacement::AtStart ? "nulls-first" + : "nulls-last"); + std::vector invalid = + null_placement == compute::NullPlacement::AtStart + ? std::vector{ExecBatchFromJSON({int64()}, "[[0]]"), + ExecBatchFromJSON({int64()}, "[[null]]")} + : std::vector{ExecBatchFromJSON({int64()}, "[[null]]"), + ExecBatchFromJSON({int64()}, "[[0]]")}; + Declaration join{"asofjoin", + {make_source(left_schema, std::move(invalid), null_placement), + make_source(right_schema, {right}, compute::NullPlacement::AtEnd)}, + GetRepeatedOptions(2, "on", {}, 0)}; + EXPECT_RAISES_WITH_MESSAGE_THAT( + Invalid, ::testing::HasSubstr("out-of-order on-key values"), + DeclarationToExecBatches(std::move(join), /*use_threads=*/false)); + } +} + +TEST(AsofJoinTest, ValidatesAndPropagatesInputOrdering) { + auto left_schema = schema({field("on", int64()), field("label", utf8())}); + auto right_schema = schema({field("on", int64()), field("label", utf8())}); + BatchesWithSchema left_batches{{}, left_schema}; + BatchesWithSchema right_batches{{}, right_schema}; + Ordering left_ordering( + {SortKey("on", compute::SortOrder::Ascending, compute::NullPlacement::AtStart), + SortKey("label", compute::SortOrder::Descending, compute::NullPlacement::AtEnd)}); + Ordering right_ordering( + {SortKey("on", compute::SortOrder::Ascending, compute::NullPlacement::AtEnd)}); + Declaration left{ + "source", + SourceNodeOptions(left_schema, left_batches.gen(false, false), left_ordering)}; + Declaration right{ + "source", + SourceNodeOptions(right_schema, right_batches.gen(false, false), right_ordering)}; + ASSERT_OK_AND_ASSIGN(auto plan, ExecPlan::Make(*threaded_exec_context())); + Declaration join_decl{"asofjoin", + {std::move(left), std::move(right)}, + GetRepeatedOptions(2, "on", {}, 0)}; + ASSERT_OK_AND_ASSIGN(ExecNode * join, join_decl.AddToPlan(plan.get())); + + Ordering expected( + {SortKey(0, compute::SortOrder::Ascending, compute::NullPlacement::AtStart), + SortKey(1, compute::SortOrder::Descending, compute::NullPlacement::AtEnd)}); + EXPECT_TRUE(join->ordering().Equals(expected)); + + Declaration unordered_left{ + "source", SourceNodeOptions(left_schema, left_batches.gen(false, false), + Ordering::Unordered())}; + Declaration ordered_right{ + "source", + SourceNodeOptions(right_schema, right_batches.gen(false, false), right_ordering)}; + Declaration unordered_join{"asofjoin", + {std::move(unordered_left), std::move(ordered_right)}, + GetRepeatedOptions(2, "on", {}, 0)}; + EXPECT_RAISES_WITH_MESSAGE_THAT(Invalid, + ::testing::HasSubstr("has no meaningful ordering"), + unordered_join.AddToPlan(plan.get())); + + for (Ordering incompatible : {Ordering({SortKey("on", compute::SortOrder::Descending)}), + Ordering({SortKey("label"), SortKey("on")})}) { + Declaration incompatible_left{ + "source", SourceNodeOptions(left_schema, left_batches.gen(false, false), + std::move(incompatible))}; + Declaration another_ordered_right{ + "source", + SourceNodeOptions(right_schema, right_batches.gen(false, false), right_ordering)}; + Declaration incompatible_join{ + "asofjoin", + {std::move(incompatible_left), std::move(another_ordered_right)}, + GetRepeatedOptions(2, "on", {}, 0)}; + EXPECT_RAISES_WITH_MESSAGE_THAT( + Invalid, ::testing::HasSubstr("must be ordered by its ascending on-key"), + incompatible_join.AddToPlan(plan.get())); + } +} + +TEST(AsofJoinTest, SequencesJitteredInputsOnBothExecutors) { + // Keep this above AsofJoin's input high watermark. Besides checking logical order, + // this ensures later physical arrivals cannot pause a reordering input while it still + // owns the batch that closes the sequencing gap. + constexpr int64_t kRows = 128; + constexpr random::SeedType kLeftSeed = 42; + constexpr random::SeedType kRightSeed = 84; + RegisterTestNodes(); + + ASSERT_OK_AND_ASSIGN(auto on_values, gen::Step()->Generate(kRows)); + auto left_table = Table::Make(schema({field("on", int64())}), {on_values}); + auto right_table = Table::Make(schema({field("on", int64()), field("value", int64())}), + {on_values, on_values}); + + for (bool use_threads : {false, true}) { + SCOPED_TRACE(use_threads ? "threaded" : "serial"); + Declaration left = Declaration::Sequence( + {{"table_source", TableSourceNodeOptions(left_table, /*max_batch_size=*/1)}, + {"jitter", JitterNodeOptions(kLeftSeed, 4)}}); + Declaration right = Declaration::Sequence( + {{"table_source", TableSourceNodeOptions(right_table, /*max_batch_size=*/1)}, + {"jitter", JitterNodeOptions(kRightSeed, 4)}}); + QueryOptions query_options; + query_options.use_threads = use_threads; + query_options.sequence_output = false; + Declaration join{"asofjoin", + {std::move(left), std::move(right)}, + GetRepeatedOptions(2, "on", {}, 0)}; + ASSERT_OK_AND_ASSIGN(auto result, + DeclarationToExecBatches(std::move(join), query_options)); + + ASSERT_EQ(result.batches.size(), static_cast(kRows)); + for (int64_t i = 0; i < kRows; ++i) { + const ExecBatch& batch = result.batches[i]; + ASSERT_EQ(batch.index, i); + ASSERT_EQ(batch.length, 1); + EXPECT_EQ(batch.values[0].make_array()->GetScalar(0).ValueOrDie()->ToString(), + std::to_string(i)); + EXPECT_EQ(batch.values[1].make_array()->GetScalar(0).ValueOrDie()->ToString(), + std::to_string(i)); + } + } +} + +TEST(AsofJoinTest, PauseStopsUntilLeftInputFinishesThenFlushes) { + BackpressureCountingNode::Register(); + auto consumer = std::make_shared(); + BackpressureCounters left_counters; + BackpressureCounters right_counters; + auto left_schema = schema({field("on", int64())}); + auto right_schema = schema({field("on", int64()), field("value", int32())}); + std::vector left_batches = { + ExecBatchFromJSON({int64()}, "[[1]]"), + ExecBatchFromJSON({int64()}, "[[2]]"), + ExecBatchFromJSON({int64()}, "[[3]]"), + }; + ExecBatch right_batch = + ExecBatchFromJSON({int64(), int32()}, "[[1, 10], [2, 20], [3, 30]]"); + + PushGenerator> left_generator; + Declaration left_source{ + "source", SourceNodeOptions(left_schema, left_generator, Ordering::Implicit())}; + Declaration left{BackpressureCountingNode::kFactoryName, + {std::move(left_source)}, + BackpressureCountingNodeOptions(&left_counters)}; + Declaration right_source{"exec_batch_source", + ExecBatchSourceNodeOptions(right_schema, {right_batch})}; + Declaration right{BackpressureCountingNode::kFactoryName, + {std::move(right_source)}, + BackpressureCountingNodeOptions(&right_counters)}; + Declaration join{"asofjoin", + {std::move(left), std::move(right)}, + GetRepeatedOptions(2, "on", {}, 0)}; + Declaration sink{ + "consuming_sink", + {std::move(join)}, + ConsumingSinkNodeOptions(consumer, /*names=*/{}, /*sequence_output=*/false)}; + + ASSERT_OK_AND_ASSIGN(auto plan, ExecPlan::Make(*threaded_exec_context())); + ASSERT_OK(sink.AddToPlan(plan.get())); + ASSERT_OK(plan->Validate()); + plan->StartProducing(); + left_generator.producer().Push(left_batches[0]); + BusyWait(10.0, [&] { return consumer->is_paused(); }); + + // The open source owns a scheduler task. Additional batches may arrive after the + // downstream pause, but no new left generation starts before the source reaches EOS. + left_generator.producer().Push(left_batches[1]); + left_generator.producer().Push(left_batches[2]); + + BusyWait(10.0, [&] { + return left_counters.batch_count.load() == 3 && + right_counters.batch_count.load() == 1; + }); + EXPECT_EQ(left_counters.batch_count.load(), 3); + EXPECT_EQ(right_counters.batch_count.load(), 1); + + // Downstream pause controls left-batch activation, not input production. With + // fewer than the queue high watermark, neither input should be paused directly. + EXPECT_EQ(left_counters.pause_count.load(), 0); + EXPECT_EQ(right_counters.pause_count.load(), 0); + + arrow::internal::GetCpuThreadPool()->WaitForIdle(); + EXPECT_TRUE(consumer->is_paused()); + EXPECT_EQ(consumer->batches_received(), 1); + EXPECT_FALSE(consumer->finished().is_finished()); + EXPECT_FALSE(plan->finished().is_finished()); + + left_generator.producer().Push(IterationEnd>()); + ASSERT_FINISHES_OK(consumer->finished()); + ASSERT_FINISHES_OK(plan->finished()); + EXPECT_EQ(consumer->batches_received(), 3); + + // A late resume after flushing is harmless because the join is already terminal. + consumer->Resume(); + EXPECT_FALSE(consumer->is_paused()); +} + +TEST(AsofJoinTest, StopWhilePausedAtLeftBatchBoundary) { + auto consumer = std::make_shared(); + auto left_schema = schema({field("on", int64())}); + auto right_schema = schema({field("on", int64()), field("value", int32())}); + ExecBatch left_batch = ExecBatchFromJSON({int64()}, "[[1]]"); + ExecBatch right_batch = ExecBatchFromJSON({int64(), int32()}, "[[1, 10]]"); + + PushGenerator> left_generator; + Declaration left{"source", + SourceNodeOptions(left_schema, left_generator, Ordering::Implicit())}; + Declaration right{"exec_batch_source", + ExecBatchSourceNodeOptions(right_schema, {right_batch})}; + Declaration join{"asofjoin", + {std::move(left), std::move(right)}, + GetRepeatedOptions(2, "on", {}, 0)}; + Declaration sink{ + "consuming_sink", + {std::move(join)}, + ConsumingSinkNodeOptions(consumer, /*names=*/{}, /*sequence_output=*/false)}; + + ASSERT_OK_AND_ASSIGN(auto plan, ExecPlan::Make(*threaded_exec_context())); + ASSERT_OK(sink.AddToPlan(plan.get())); + ASSERT_OK(plan->Validate()); + plan->StartProducing(); + left_generator.producer().Push(std::move(left_batch)); + BusyWait(10.0, [&] { return consumer->is_paused(); }); + EXPECT_TRUE(consumer->is_paused()); + arrow::internal::GetCpuThreadPool()->WaitForIdle(); + EXPECT_EQ(consumer->batches_received(), 1); + EXPECT_FALSE(plan->finished().is_finished()); + + plan->StopProducing(); + left_generator.producer().Push(IterationEnd>()); + ASSERT_TRUE(plan->finished().Wait(kDefaultAssertFinishesWaitSeconds)); + ASSERT_TRUE(plan->finished().status().IsCancelled()); +} + +TEST(AsofJoinTest, PauseAllowsActiveLeftBatchToFinish) { + auto consumer = std::make_shared(/*pause_after_first=*/false); + auto left_schema = schema({field("on", int64())}); + auto right_schema = schema({field("on", int64()), field("value", int32())}); + ExecBatch left_batch = ExecBatchFromJSON({int64()}, "[[1]]"); + ExecBatch right_batch = ExecBatchFromJSON({int64(), int32()}, "[[1, 10], [2, 20]]"); + + Future> right_ready = Future>::Make(); + auto generator_calls = std::make_shared>(0); + AsyncGenerator> delayed_right = [right_ready, + generator_calls]() mutable { + if (generator_calls->fetch_add(1) == 0) { + return right_ready; + } + return Future>::MakeFinished(std::nullopt); + }; + + Declaration left{"exec_batch_source", + ExecBatchSourceNodeOptions(left_schema, {left_batch})}; + Declaration right{"source", SourceNodeOptions(right_schema, std::move(delayed_right), + Ordering::Implicit())}; + Declaration join{"asofjoin", + {std::move(left), std::move(right)}, + GetRepeatedOptions(2, "on", {}, 0)}; + Declaration sink{ + "consuming_sink", + {std::move(join)}, + ConsumingSinkNodeOptions(consumer, /*names=*/{}, /*sequence_output=*/false)}; + + ASSERT_OK_AND_ASSIGN(auto plan, ExecPlan::Make(*threaded_exec_context())); + ASSERT_OK(sink.AddToPlan(plan.get())); + ASSERT_OK(plan->Validate()); + plan->StartProducing(); + + // The source's unresolved future is tracked by the plan scheduler. Once the CPU + // executor is idle, the left lane is active and waiting for this RHS batch. + arrow::internal::GetCpuThreadPool()->WaitForIdle(); + ASSERT_TRUE(consumer->is_initialized()); + ASSERT_EQ(consumer->batches_received(), 0); + consumer->Pause(); + right_ready.MarkFinished(std::optional(std::move(right_batch))); + + BusyWait(10.0, [&] { return consumer->batches_received() == 1; }); + if (consumer->batches_received() != 1) { + // Ensure a failing test still drains the plan cleanly. + consumer->Resume(); + ASSERT_FINISHES_OK(consumer->finished()); + ASSERT_FINISHES_OK(plan->finished()); + FAIL() << "an already-active left batch remained stranded behind pause"; + } + ASSERT_FINISHES_OK(consumer->finished()); + ASSERT_FINISHES_OK(plan->finished()); + consumer->Resume(); + EXPECT_EQ(consumer->batches_received(), 1); +} + // Reproduction of GH-40675: A logical race between Process() and Push() that can be more // easily observed with single small batch. TEST(AsofJoinTest, RhsEmptinessRace) { diff --git a/cpp/src/arrow/acero/exec_plan.h b/cpp/src/arrow/acero/exec_plan.h index dba6c64ddc83..8a66eab4f193 100644 --- a/cpp/src/arrow/acero/exec_plan.h +++ b/cpp/src/arrow/acero/exec_plan.h @@ -193,9 +193,10 @@ class ARROW_ACERO_EXPORT ExecNode { /// new ordering based on the hash keys). /// /// Some nodes will require an ordering. For example, a fetch node or an - /// asof join node will only function if the input data is ordered (for fetch - /// it is enough to be implicitly ordered. For an asof join the ordering must - /// be explicit and compatible with the on key.) + /// asof join node will only function if the input data is ordered. For fetch, + /// implicit ordering is sufficient. An asof join accepts implicit ordering after + /// validating the on key at runtime; explicit ordering must be compatible with the + /// on key. /// /// Nodes that maintain ordering should be careful to avoid introducing gaps /// in the batch index. This may require emitting empty batches in order to diff --git a/cpp/src/arrow/acero/options.h b/cpp/src/arrow/acero/options.h index 827e9ea775d7..d273452188aa 100644 --- a/cpp/src/arrow/acero/options.h +++ b/cpp/src/arrow/acero/options.h @@ -687,7 +687,10 @@ class ARROW_ACERO_EXPORT HashJoinNodeOptions : public ExecNodeOptions { /// Note, this API is experimental and will change in the future /// /// This node takes one left table and any number of right tables, and asof joins them -/// together. Batches produced by each input must be ordered by the "on" key. +/// together. Batches produced by each input must have a meaningful ordering, and +/// non-null "on" values must be ordered ascending. Explicit input ordering must have the +/// "on" key as its leading ascending sort key and follows that key's null placement. +/// Null placement is ignored for implicitly ordered input. /// This node will output one row for each row in the left table. class ARROW_ACERO_EXPORT AsofJoinNodeOptions : public ExecNodeOptions { public: @@ -699,23 +702,42 @@ class ARROW_ACERO_EXPORT AsofJoinNodeOptions : public ExecNodeOptions { struct Keys { /// \brief "on" key for the join. /// - /// The input table must be sorted by the "on" key. Must be a single field of a common - /// type. An inexact match is used on the "on" key, i.e. a row is considered a - /// match if and only if `right.on - left.on` is in the range - /// `[min(0, tolerance), max(0, tolerance)]`. - /// Currently, the "on" key must be of an integer, date, or timestamp type. + /// Non-null values must be sorted ascending. For explicitly ordered input, this must + /// be the leading ascending sort key and its declared null placement must be obeyed. + /// Null placement is ignored for implicitly ordered input. + /// Must be a single field of a common type. An inexact match is used on the "on" key, + /// i.e. a row is eligible if and only if `right.on - left.on` is in the configured + /// tolerance range. + /// Currently, the "on" key must be of an integer, date, time, or timestamp type. FieldRef on_key; /// \brief "by" key for the join. /// /// Each input table must have each field of the "by" key. Exact equality is used for /// each field of the "by" key. - /// Currently, each field of the "by" key must be of an integer, date, timestamp, or - /// base-binary type. + /// Currently, each field of the "by" key must be of a boolean, integer, date, time, + /// timestamp, string, binary, fixed-size binary, or decimal type. Dictionary-encoded + /// fields are not supported. std::vector by_key; }; + /// \brief Inclusive tolerance range for inexact "on" key matching. + /// + /// A right row is eligible when `right.on - left.on` is in `[lower, upper]`. + struct ToleranceRange { + int64_t lower; + int64_t upper; + }; + AsofJoinNodeOptions(std::vector input_keys, int64_t tolerance) - : input_keys(std::move(input_keys)), tolerance(tolerance) {} + : AsofJoinNodeOptions(std::move(input_keys), + ToleranceRange{tolerance < 0 ? tolerance : 0, + tolerance > 0 ? tolerance : 0}) {} + + AsofJoinNodeOptions(std::vector input_keys, ToleranceRange tolerance, + bool prefer_earlier_on_tie = true) + : input_keys(std::move(input_keys)), + tolerance(std::move(tolerance)), + prefer_earlier_on_tie(prefer_earlier_on_tie) {} /// \brief AsofJoin keys per input table. At least two keys must be given. The first key /// corresponds to a left table and all other keys correspond to right tables for the @@ -723,18 +745,16 @@ class ARROW_ACERO_EXPORT AsofJoinNodeOptions : public ExecNodeOptions { /// /// \see `Keys` for details. std::vector input_keys; - /// \brief Tolerance for inexact "on" key matching. A right row is considered a match - /// with a left row if `right.on - left.on` is in the range - /// `[min(0, tolerance), max(0, tolerance)]`. `tolerance` may be: - /// - negative, in which case a past-as-of-join occurs (match iff - /// `tolerance <= right.on - left.on <= 0`); - /// - or positive, in which case a future-as-of-join occurs (match iff - /// `0 <= right.on - left.on <= tolerance`); - /// - or zero, in which case an exact-as-of-join occurs (match iff - /// `right.on == left.on`). - /// - /// The tolerance is interpreted in the same units as the "on" key. - int64_t tolerance; + /// \brief Tolerance range for inexact "on" key matching. + /// + /// The bounds are interpreted in the same units as the "on" key. `lower` must not be + /// greater than `upper`. + ToleranceRange tolerance; + /// \brief Which side to prefer when eligible rows are equally close. + /// + /// If two rows on opposite sides of the left row are equally close, choose the row + /// with the smaller "on" value when true and the larger "on" value when false. + bool prefer_earlier_on_tie; }; /// \brief a node which select top_k/bottom_k rows passed through it diff --git a/cpp/src/arrow/acero/time_series_util.cc b/cpp/src/arrow/acero/time_series_util.cc index 60f9044d7bf7..9fbf54666373 100644 --- a/cpp/src/arrow/acero/time_series_util.cc +++ b/cpp/src/arrow/acero/time_series_util.cc @@ -22,13 +22,6 @@ namespace arrow::acero { -template ::value, bool>> -inline uint64_t NormalizeTime(T t) { - uint64_t bias = - std::is_signed::value ? static_cast(1) << (8 * sizeof(T) - 1) : 0; - return t < 0 ? static_cast(t + bias) : static_cast(t); -} - uint64_t GetTime(const RecordBatch* batch, Type::type time_type, int col, uint64_t row) { #define LATEST_VAL_CASE(id, val) \ case Type::id: { \ diff --git a/cpp/src/arrow/acero/time_series_util.h b/cpp/src/arrow/acero/time_series_util.h index 97707f43bf20..852a647002ad 100644 --- a/cpp/src/arrow/acero/time_series_util.h +++ b/cpp/src/arrow/acero/time_series_util.h @@ -17,14 +17,25 @@ #pragma once +#include +#include + #include "arrow/record_batch.h" #include "arrow/type_traits.h" namespace arrow::acero { // normalize the value to unsigned 64-bits while preserving ordering of values -template ::value, bool> = true> -uint64_t NormalizeTime(T t); +template && !std::is_same_v, bool> = true> +uint64_t NormalizeTime(T t) { + using U = std::make_unsigned_t; + U normalized = static_cast(t); + if constexpr (std::is_signed_v) { + normalized ^= U{1} << (std::numeric_limits::digits - 1); + } + return static_cast(normalized); +} uint64_t GetTime(const RecordBatch* batch, Type::type time_type, int col, uint64_t row); diff --git a/docs/source/developers/cpp/acero.rst b/docs/source/developers/cpp/acero.rst index d024a2d03d60..f9cec581e1fb 100644 --- a/docs/source/developers/cpp/acero.rst +++ b/docs/source/developers/cpp/acero.rst @@ -181,10 +181,10 @@ If any external resources are used then cleanup should happen as part of this ca Examples ^^^^^^^^ -* The ``asofjoin`` node has a dedicated processing thread the communicates with the main Acero threads - using a queue. When ``StopProducing`` is called the node inserts a poison pill into the queue. This - tells the processing thread to stop immediately. Once the processing thread stops it marks its external - task (described below) as completed which allows the plan to finish. +* The ``asofjoin`` node uses tracked lane tasks to process its inputs. When + ``StopProducing`` is called, the node marks its coordinator and lanes as stopped and + stops its inputs. Tasks that were already scheduled observe the terminal state and + return without publishing more output. * The ``fetch`` node, in ``InputReceived``, may decide that it has all the data it needs. It can then call ``StopProducing`` on its input. @@ -403,10 +403,9 @@ source nodes will usually schedule tasks during the call to StartProducing. Pip when they have accumulated all the data they need. Once all tasks in a plan are finished then the plan is considered done. -Some nodes use external threads. These threads must be registered as external tasks using the BeginExternalTask method. -For example, the asof join node uses a dedicated processing thread to achieve serial execution. This dedicated thread -is registered as an external task. External tasks should be avoided where possible because they require careful -handling to avoid deadlock in error situations. +Some nodes use external threads. These threads must be registered as external tasks +using the BeginExternalTask method. External tasks should be avoided where possible +because they require careful handling to avoid deadlock in error situations. Ordered Execution ================= @@ -536,8 +535,9 @@ If you need to run something in parallel then you should use thread tasks and no * This makes it possible to run without threads (sometimes users are doing their own threading and sometimes we need to run in thread-restricted environments like emscripten) -Note: we do not always follow this advice currently. There is a dedicated process thread in the asof join -node. Dedicated threads are "ok" for experimental use but we'd like to migrate away from them. +The ``asofjoin`` node follows this model by sequencing each input and scheduling tracked +lane tasks. The same implementation therefore runs on either a serial executor or a +thread pool without owning a dedicated thread. Don't Block on CPU Threads -------------------------- diff --git a/python/pyarrow/_acero.pyx b/python/pyarrow/_acero.pyx index e27509469ef5..e6d280fe9b39 100644 --- a/python/pyarrow/_acero.pyx +++ b/python/pyarrow/_acero.pyx @@ -436,13 +436,15 @@ class HashJoinNodeOptions(_HashJoinNodeOptions): cdef class _AsofJoinNodeOptions(ExecNodeOptions): - def _set_options(self, left_on, left_by, right_on, right_by, tolerance): + def _set_options(self, left_on, left_by, right_on, right_by, tolerance, + prefer_earlier_on_tie): cdef: vector[CFieldRef] c_left_by vector[CFieldRef] c_right_by CAsofJoinKeys c_left_keys CAsofJoinKeys c_right_keys vector[CAsofJoinKeys] c_input_keys + CAsofJoinToleranceRange c_tolerance # Prepare left AsofJoinNodeOption::Keys if not isinstance(left_by, (list, tuple)): @@ -466,10 +468,19 @@ cdef class _AsofJoinNodeOptions(ExecNodeOptions): c_input_keys.push_back(c_right_keys) + if isinstance(tolerance, (list, tuple)): + if len(tolerance) != 2: + raise ValueError("tolerance range must contain lower and upper bounds") + c_tolerance.lower = tolerance[0] + c_tolerance.upper = tolerance[1] + else: + c_tolerance.lower = tolerance if tolerance < 0 else 0 + c_tolerance.upper = tolerance if tolerance > 0 else 0 self.wrapped.reset( new CAsofJoinNodeOptions( c_input_keys, - tolerance, + c_tolerance, + prefer_earlier_on_tie, ) ) @@ -486,14 +497,13 @@ class AsofJoinNodeOptions(_AsofJoinNodeOptions): The left key on which the join operation should be performed. Can be a string column name or a field expression. - An inexact match is used on the "on" key, i.e. a row is considered a - match if and only if ``right.on - left.on`` is in the range - ``[min(0, tolerance), max(0, tolerance)]``. + An inexact match is used on the "on" key. A row is eligible when + ``right.on - left.on`` is in the configured tolerance range. The input dataset must be sorted by the "on" key. Must be a single field of a common type. - Currently, the "on" key must be an integer, date, or timestamp type. + Currently, the "on" key must be an integer, date, time, or timestamp type. left_by: str, Expression or list The left keys on which the join operation should be performed. Exact equality is used for each field of the "by" keys. @@ -505,13 +515,20 @@ class AsofJoinNodeOptions(_AsofJoinNodeOptions): right_by: str, Expression or list The right keys on which the join operation should be performed. See `left_by` for details. - tolerance : int - The tolerance to use for the asof join. The tolerance is interpreted in - the same units as the "on" key. + tolerance : int or 2-tuple of int + The tolerance to use for the asof join. A scalar preserves the legacy + directional behavior. A tuple gives the inclusive lower and upper bounds + for ``right.on - left.on``. Bounds are interpreted in the same units as + the "on" key. + prefer_earlier_on_tie : bool, default True + If two eligible rows are equally close, choose the row with the smaller + "on" value when true and the larger "on" value when false. """ - def __init__(self, left_on, left_by, right_on, right_by, tolerance): - self._set_options(left_on, left_by, right_on, right_by, tolerance) + def __init__(self, left_on, left_by, right_on, right_by, tolerance, + prefer_earlier_on_tie=True): + self._set_options(left_on, left_by, right_on, right_by, tolerance, + prefer_earlier_on_tie) cdef class Declaration(_Weakrefable): diff --git a/python/pyarrow/_dataset.pyx b/python/pyarrow/_dataset.pyx index d40614a61fc0..020f10743902 100644 --- a/python/pyarrow/_dataset.pyx +++ b/python/pyarrow/_dataset.pyx @@ -924,7 +924,8 @@ cdef class Dataset(_Weakrefable): output_type=InMemoryDataset ) - def join_asof(self, right_dataset, on, by, tolerance, right_on=None, right_by=None): + def join_asof(self, right_dataset, on, by, tolerance, right_on=None, right_by=None, + prefer_earlier_on_tie=True): """ Perform an asof join between this dataset and another one. @@ -946,22 +947,21 @@ cdef class Dataset(_Weakrefable): The column from current dataset that should be used as the "on" key of the join operation left side. - An inexact match is used on the "on" key, i.e. a row is considered a - match if and only if ``right.on - left.on`` is in the range - ``[min(0, tolerance), max(0, tolerance)]``. + An inexact match is used on the "on" key. A row is eligible when + ``right.on - left.on`` is in the configured tolerance range. The input table must be sorted by the "on" key. Must be a single field of a common type. - Currently, the "on" key must be an integer, date, or timestamp type. + Currently, the "on" key must be an integer, date, time, or timestamp type. by : str or list[str] The columns from current dataset that should be used as the keys of the join operation left side. The join operation is then done only for the matches in these columns. - tolerance : int + tolerance : int or 2-tuple of int The tolerance for inexact "on" key matching. A right row is considered - a match with a left row if ``right.on - left.on`` is in the range - ``[min(0, tolerance), max(0, tolerance)]``. ``tolerance`` may be: + eligible when ``right.on - left.on`` is in the configured inclusive + range. A scalar ``tolerance`` preserves the directional behavior: - negative, in which case a past-as-of-join occurs (match iff ``tolerance <= right.on - left.on <= 0``); @@ -970,6 +970,10 @@ cdef class Dataset(_Weakrefable): - or zero, in which case an exact-as-of-join occurs (match iff ``right.on == left.on``). + A 2-tuple gives the lower and upper bounds directly. For example, + ``(-10, -1)`` performs a backward join that excludes exact matches, + while ``(-5, 10)`` searches on both sides. + The tolerance is interpreted in the same units as the "on" key. right_on : str or list[str], default None The columns from the right_dataset that should be used as the on key @@ -979,6 +983,9 @@ cdef class Dataset(_Weakrefable): The columns from the right_dataset that should be used as by keys on the join operation right side. When ``None`` use the same key names as the left dataset. + prefer_earlier_on_tie : bool, default True + If two eligible rows are equally close, choose the row with the smaller + "on" value when true and the larger "on" value when false. Returns ------- @@ -990,7 +997,8 @@ cdef class Dataset(_Weakrefable): right_by = by return _pac()._perform_join_asof(self, on, by, right_dataset, right_on, right_by, - tolerance, output_type=InMemoryDataset) + tolerance, prefer_earlier_on_tie, + output_type=InMemoryDataset) cdef class InMemoryDataset(Dataset): diff --git a/python/pyarrow/acero.py b/python/pyarrow/acero.py index 55f03a8bb75f..14b08e22780d 100644 --- a/python/pyarrow/acero.py +++ b/python/pyarrow/acero.py @@ -263,7 +263,7 @@ def _perform_join(join_type, left_operand, left_keys, def _perform_join_asof(left_operand, left_on, left_by, right_operand, right_on, right_by, - tolerance, use_threads=True, + tolerance, prefer_earlier_on_tie=True, use_threads=True, output_type=Table): """ Perform asof join of two tables or datasets. @@ -284,9 +284,13 @@ def _perform_join_asof(left_operand, left_on, left_by, The right key (or keys) on which the join operation should be performed. right_by: str or list[str] The right key (or keys) on which the join operation should be performed. - tolerance : int - The tolerance to use for the asof join. The tolerance is interpreted in - the same units as the "on" key. + tolerance : int or 2-tuple of int + The tolerance to use for the asof join. A tuple gives the inclusive lower + and upper bounds for ``right.on - left.on``. Bounds are interpreted in the + same units as the "on" key. + prefer_earlier_on_tie : bool, default True + If two eligible rows are equally close, choose the row with the smaller + "on" value when true and the larger "on" value when false. output_type: Table or InMemoryDataset The output type for the exec plan result. @@ -336,7 +340,8 @@ def _perform_join_asof(left_operand, left_on, left_by, ) join_opts = AsofJoinNodeOptions( - left_on, left_by, right_on, right_by, tolerance + left_on, left_by, right_on, right_by, tolerance, + prefer_earlier_on_tie ) decl = Declaration( "asofjoin", options=join_opts, inputs=[left_source, right_source] diff --git a/python/pyarrow/includes/libarrow_acero.pxd b/python/pyarrow/includes/libarrow_acero.pxd index 7e81a393682c..dbd9498e22b2 100644 --- a/python/pyarrow/includes/libarrow_acero.pxd +++ b/python/pyarrow/includes/libarrow_acero.pxd @@ -83,8 +83,15 @@ cdef extern from "arrow/acero/options.h" namespace "arrow::acero" nogil: CFieldRef on_key vector[CFieldRef] by_key + cdef struct CAsofJoinToleranceRange "arrow::acero::AsofJoinNodeOptions::ToleranceRange": + int64_t lower + int64_t upper + cdef cppclass CAsofJoinNodeOptions "arrow::acero::AsofJoinNodeOptions"(CExecNodeOptions): CAsofJoinNodeOptions(vector[CAsofJoinKeys] keys, int64_t tolerance) + CAsofJoinNodeOptions(vector[CAsofJoinKeys] keys, + CAsofJoinToleranceRange tolerance, + c_bool prefer_earlier_on_tie) cdef extern from "arrow/acero/exec_plan.h" namespace "arrow::acero" nogil: diff --git a/python/pyarrow/table.pxi b/python/pyarrow/table.pxi index 1abe4235c411..7f63e3521348 100644 --- a/python/pyarrow/table.pxi +++ b/python/pyarrow/table.pxi @@ -5791,7 +5791,8 @@ cdef class Table(_Tabular): filter_expression=filter_expression, ) - def join_asof(self, right_table, on, by, tolerance, right_on=None, right_by=None): + def join_asof(self, right_table, on, by, tolerance, right_on=None, right_by=None, + prefer_earlier_on_tie=True): """ Perform an asof join between this table and another one. @@ -5813,22 +5814,21 @@ cdef class Table(_Tabular): The column from current table that should be used as the "on" key of the join operation left side. - An inexact match is used on the "on" key, i.e. a row is considered a - match if and only if ``right.on - left.on`` is in the range - ``[min(0, tolerance), max(0, tolerance)]``. + An inexact match is used on the "on" key. A row is eligible when + ``right.on - left.on`` is in the configured tolerance range. The input dataset must be sorted by the "on" key. Must be a single field of a common type. - Currently, the "on" key must be an integer, date, or timestamp type. + Currently, the "on" key must be an integer, date, time, or timestamp type. by : str or list[str] The columns from current table that should be used as the keys of the join operation left side. The join operation is then done only for the matches in these columns. - tolerance : int + tolerance : int or 2-tuple of int The tolerance for inexact "on" key matching. A right row is considered - a match with a left row if ``right.on - left.on`` is in the range - ``[min(0, tolerance), max(0, tolerance)]``. ``tolerance`` may be: + eligible when ``right.on - left.on`` is in the configured inclusive + range. A scalar ``tolerance`` preserves the directional behavior: - negative, in which case a past-as-of-join occurs (match iff ``tolerance <= right.on - left.on <= 0``); @@ -5837,6 +5837,10 @@ cdef class Table(_Tabular): - or zero, in which case an exact-as-of-join occurs (match iff ``right.on == left.on``). + A 2-tuple gives the lower and upper bounds directly. For example, + ``(-10, -1)`` performs a backward join that excludes exact matches, + while ``(-5, 10)`` searches on both sides. + The tolerance is interpreted in the same units as the "on" key. right_on : str or list[str], default None The columns from the right_table that should be used as the on key @@ -5846,6 +5850,9 @@ cdef class Table(_Tabular): The columns from the right_table that should be used as keys on the join operation right side. When ``None`` use the same key names as the left table. + prefer_earlier_on_tie : bool, default True + If two eligible rows are equally close, choose the row with the smaller + "on" value when true and the larger "on" value when false. Returns ------- @@ -5880,7 +5887,8 @@ cdef class Table(_Tabular): right_by = by return _pac()._perform_join_asof(self, on, by, right_table, right_on, right_by, - tolerance, output_type=Table) + tolerance, prefer_earlier_on_tie, + output_type=Table) def __arrow_c_stream__(self, requested_schema=None): """ diff --git a/python/pyarrow/tests/test_acero.py b/python/pyarrow/tests/test_acero.py index 6e471d61199f..0f0ceef5a7a4 100644 --- a/python/pyarrow/tests/test_acero.py +++ b/python/pyarrow/tests/test_acero.py @@ -532,7 +532,7 @@ def test_hash_join_with_residual_filter(): def test_asof_join(): left = pa.table({'key': [1, 2, 3], 'ts': [1, 1, 1], 'a': [4, 5, 6]}) left_source = Declaration("table_source", options=TableSourceNodeOptions(left)) - right = pa.table({'key': [2, 3, 4], 'ts': [2, 5, 2], 'b': [4, 5, 6]}) + right = pa.table({'key': [2, 4, 3], 'ts': [2, 2, 5], 'b': [4, 6, 5]}) right_source = Declaration("table_source", options=TableSourceNodeOptions(right)) # asof join diff --git a/python/pyarrow/tests/test_dataset.py b/python/pyarrow/tests/test_dataset.py index 09d7cfb9d9dd..d4d7f5f5f4a4 100644 --- a/python/pyarrow/tests/test_dataset.py +++ b/python/pyarrow/tests/test_dataset.py @@ -5401,6 +5401,32 @@ def test_dataset_join_asof(tempdir): }) +@pytest.mark.dataset +def test_dataset_join_asof_tolerance_range(tempdir): + left = pa.table({ + "left": [10, 20, 30, 40], + "key": [1, 1, 1, 1], + }) + right = pa.table({ + "right": [9, 12, 30, 41], + "key": [1, 1, 1, 1], + "value": [1, 2, 3, 4], + }) + ds.write_dataset(left, tempdir / "left", format="ipc") + ds.write_dataset(right, tempdir / "right", format="ipc") + + result = ds.dataset(tempdir / "left", format="ipc").join_asof( + ds.dataset(tempdir / "right", format="ipc"), + on="left", by="key", tolerance=(-10, -1), + right_on="right", right_by="key", + ) + assert result.to_table() == pa.table({ + "left": [10, 20, 30, 40], + "key": [1, 1, 1, 1], + "value": [1, 2, None, 3], + }) + + @pytest.mark.dataset def test_dataset_join_asof_multiple_by(tempdir): t1 = pa.table({ diff --git a/python/pyarrow/tests/test_table.py b/python/pyarrow/tests/test_table.py index bf6e5773ddf1..5d02ca5b6668 100644 --- a/python/pyarrow/tests/test_table.py +++ b/python/pyarrow/tests/test_table.py @@ -3287,6 +3287,60 @@ def test_table_join_asof(): }) +@pytest.mark.dataset +def test_table_join_asof_tolerance_range(): + left = pa.table({ + "left": [10, 20, 30, 40], + "key": [1, 1, 1, 1], + }) + right = pa.table({ + "right": [9, 12, 30, 41], + "key": [1, 1, 1, 1], + "value": [1, 2, 3, 4], + }) + + result = left.join_asof( + right, on="left", by="key", tolerance=(-10, -1), + right_on="right", right_by="key", + ) + assert result == pa.table({ + "left": [10, 20, 30, 40], + "key": [1, 1, 1, 1], + "value": [1, 2, None, 3], + }) + + tie_left = pa.table({"on": [10], "key": [1]}) + tie_right = pa.table({ + "on": [8, 12], + "key": [1, 1], + "value": [8, 12], + }) + assert tie_left.join_asof( + tie_right, on="on", by="key", tolerance=(-2, 2) + )["value"].to_pylist() == [8] + assert tie_left.join_asof( + tie_right, on="on", by="key", tolerance=(-2, 2), + prefer_earlier_on_tie=False, + )["value"].to_pylist() == [12] + + with pytest.raises(pa.ArrowInvalid, match="lower bound must not exceed"): + left.join_asof( + right, on="left", by="key", tolerance=(1, -1), + right_on="right", right_by="key", + ) + + +@pytest.mark.dataset +def test_table_join_asof_null_on_keys(): + left = pa.table({"time": pa.array([None], type=pa.int64())}) + right = pa.table({"time": [0], "value": [True]}) + + assert left.join_asof(right, on="time", by=[], tolerance=0) == pa.table({ + "time": pa.array([None], type=pa.int64()), + "value": pa.array([None], type=pa.bool_()), + }) + + @pytest.mark.dataset def test_table_join_asof_multiple_by(): t1 = pa.table({