You are viewing a plain text version of this content. The canonical link for it is here.
Posted to github@arrow.apache.org by GitBox <gi...@apache.org> on 2022/05/26 04:12:03 UTC
[GitHub] [arrow] westonpace commented on a diff in pull request #13028: ARROW-16083: [C++] Implement AsofJoin execution node
westonpace commented on code in PR #13028:
URL: https://github.com/apache/arrow/pull/13028#discussion_r882244204
##########
cpp/src/arrow/compute/exec/asof_join.h:
##########
@@ -0,0 +1,42 @@
+// Licensed to the Apache Software Foundation (ASF) under one
Review Comment:
Shall we get rid of this file and move these declarations into the top of `asof_join_node.cc`? Or, if you are planning on building other utilities around this header file can you move the definition of `AsofJoinSchema::MakeOutputSchema` into an `asof_join.cc`?
##########
cpp/src/arrow/compute/exec/asof_join.h:
##########
@@ -0,0 +1,42 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+#pragma once
+
+#include <arrow/compute/api.h>
+#include <arrow/compute/exec/exec_plan.h>
+#include "arrow/compute/exec/options.h"
+#include "arrow/compute/exec/schema_util.h"
+
+namespace arrow {
+namespace compute {
+
+typedef int32_t KeyType;
+
+// Maximum number of tables that can be joined
+#define MAX_JOIN_TABLES 64
Review Comment:
```suggestion
constexpr int kMaxJoinTables = 64
```
Macros ignore namespaces and are prone to conflicts. We try to avoid them when possible.
##########
cpp/src/arrow/compute/exec/asof_join_node.cc:
##########
@@ -0,0 +1,800 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+#include <iostream>
+#include <unordered_map>
+
+#include <arrow/api.h>
+#include <arrow/compute/api.h>
+#include <arrow/util/optional.h>
+#include "arrow/compute/exec/asof_join.h"
+#include "arrow/compute/exec/exec_plan.h"
+#include "arrow/compute/exec/options.h"
+#include "arrow/compute/exec/schema_util.h"
+#include "arrow/compute/exec/util.h"
+#include "arrow/result.h"
+#include "arrow/status.h"
+#include "arrow/util/checked_cast.h"
+#include "arrow/util/future.h"
+#include "arrow/util/make_unique.h"
+
+#include <condition_variable>
+#include <mutex>
+#include <thread>
+
+namespace arrow {
+namespace compute {
+
+/**
+ * Simple implementation for an unbound concurrent queue
+ */
+template <class T>
+class ConcurrentQueue {
+ public:
+ T pop() {
Review Comment:
```suggestion
T Pop() {
```
##########
cpp/src/arrow/compute/exec/asof_join_node.cc:
##########
@@ -0,0 +1,800 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+#include <iostream>
+#include <unordered_map>
+
+#include <arrow/api.h>
+#include <arrow/compute/api.h>
+#include <arrow/util/optional.h>
+#include "arrow/compute/exec/asof_join.h"
+#include "arrow/compute/exec/exec_plan.h"
+#include "arrow/compute/exec/options.h"
+#include "arrow/compute/exec/schema_util.h"
+#include "arrow/compute/exec/util.h"
+#include "arrow/result.h"
+#include "arrow/status.h"
+#include "arrow/util/checked_cast.h"
+#include "arrow/util/future.h"
+#include "arrow/util/make_unique.h"
+
+#include <condition_variable>
+#include <mutex>
+#include <thread>
+
+namespace arrow {
+namespace compute {
+
+/**
+ * Simple implementation for an unbound concurrent queue
+ */
+template <class T>
+class ConcurrentQueue {
+ public:
+ T pop() {
+ std::unique_lock<std::mutex> lock(mutex_);
+ cond_.wait(lock, [&] { return !queue_.empty(); });
+ auto item = queue_.front();
+ queue_.pop();
+ return item;
+ }
+
+ void push(const T& item) {
+ std::unique_lock<std::mutex> lock(mutex_);
+ queue_.push(item);
+ cond_.notify_one();
+ }
+
+ util::optional<T> try_pop() {
+ // Try to pop the oldest value from the queue (or return nullopt if none)
+ std::unique_lock<std::mutex> lock(mutex_);
+ if (queue_.empty()) {
+ return util::nullopt;
+ } else {
+ auto item = queue_.front();
+ queue_.pop();
+ return item;
+ }
+ }
+
+ bool empty() const {
+ std::unique_lock<std::mutex> lock(mutex_);
+ return queue_.empty();
+ }
+
+ // Un-synchronized access to front
+ // For this to be "safe":
+ // 1) the caller logically guarantees that queue is not empty
+ // 2) pop/try_pop cannot be called concurrently with this
+ const T& unsync_front() const { return queue_.front(); }
+
+ private:
+ std::queue<T> queue_;
+ mutable std::mutex mutex_;
+ std::condition_variable cond_;
+};
+
+struct MemoStore {
+ // Stores last known values for all the keys
+
+ struct Entry {
+ // Timestamp associated with the entry
+ int64_t _time;
+
+ // Batch associated with the entry (perf is probably OK for this; batches change
+ // rarely)
+ std::shared_ptr<arrow::RecordBatch> _batch;
+
+ // Row associated with the entry
+ row_index_t _row;
+ };
+
+ std::unordered_map<KeyType, Entry> _entries;
+
+ void store(const std::shared_ptr<RecordBatch>& batch, row_index_t row, int64_t time,
+ KeyType key) {
+ auto& e = _entries[key];
+ // that we can do this assignment optionally, is why we
+ // can get array 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;
+ }
+
+ util::optional<const Entry*> get_entry_for_key(KeyType key) const {
Review Comment:
```suggestion
util::optional<const Entry*> GetEntryForKey(KeyType key) const {
```
##########
cpp/src/arrow/compute/exec/asof_join_node.cc:
##########
@@ -0,0 +1,800 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+#include <iostream>
+#include <unordered_map>
+
+#include <arrow/api.h>
+#include <arrow/compute/api.h>
+#include <arrow/util/optional.h>
+#include "arrow/compute/exec/asof_join.h"
+#include "arrow/compute/exec/exec_plan.h"
+#include "arrow/compute/exec/options.h"
+#include "arrow/compute/exec/schema_util.h"
+#include "arrow/compute/exec/util.h"
+#include "arrow/result.h"
+#include "arrow/status.h"
+#include "arrow/util/checked_cast.h"
+#include "arrow/util/future.h"
+#include "arrow/util/make_unique.h"
+
+#include <condition_variable>
+#include <mutex>
+#include <thread>
+
+namespace arrow {
+namespace compute {
+
+/**
+ * Simple implementation for an unbound concurrent queue
+ */
+template <class T>
+class ConcurrentQueue {
+ public:
+ T pop() {
+ std::unique_lock<std::mutex> lock(mutex_);
+ cond_.wait(lock, [&] { return !queue_.empty(); });
+ auto item = queue_.front();
+ queue_.pop();
+ return item;
+ }
+
+ void push(const T& item) {
+ std::unique_lock<std::mutex> lock(mutex_);
+ queue_.push(item);
+ cond_.notify_one();
+ }
+
+ util::optional<T> try_pop() {
+ // Try to pop the oldest value from the queue (or return nullopt if none)
+ std::unique_lock<std::mutex> lock(mutex_);
+ if (queue_.empty()) {
+ return util::nullopt;
+ } else {
+ auto item = queue_.front();
+ queue_.pop();
+ return item;
+ }
+ }
+
+ bool empty() const {
+ std::unique_lock<std::mutex> lock(mutex_);
+ return queue_.empty();
+ }
+
+ // Un-synchronized access to front
+ // For this to be "safe":
+ // 1) the caller logically guarantees that queue is not empty
+ // 2) pop/try_pop cannot be called concurrently with this
+ const T& unsync_front() const { return queue_.front(); }
+
+ private:
+ std::queue<T> queue_;
+ mutable std::mutex mutex_;
+ std::condition_variable cond_;
+};
+
+struct MemoStore {
+ // Stores last known values for all the keys
+
+ struct Entry {
+ // Timestamp associated with the entry
+ int64_t _time;
+
+ // Batch associated with the entry (perf is probably OK for this; batches change
+ // rarely)
+ std::shared_ptr<arrow::RecordBatch> _batch;
+
+ // Row associated with the entry
+ row_index_t _row;
+ };
+
+ std::unordered_map<KeyType, Entry> _entries;
+
+ void store(const std::shared_ptr<RecordBatch>& batch, row_index_t row, int64_t time,
+ KeyType key) {
+ auto& e = _entries[key];
+ // that we can do this assignment optionally, is why we
+ // can get array 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;
+ }
+
+ util::optional<const Entry*> get_entry_for_key(KeyType key) const {
+ auto e = _entries.find(key);
+ if (_entries.end() == e) return util::nullopt;
+ return util::optional<const Entry*>(&e->second);
+ }
+
+ void remove_entries_with_lesser_time(int64_t ts) {
+ size_t dbg_size0 = _entries.size();
+ for (auto e = _entries.begin(); e != _entries.end();)
+ if (e->second._time < ts)
+ e = _entries.erase(e);
+ else
+ ++e;
+ size_t dbg_size1 = _entries.size();
+ if (dbg_size1 < dbg_size0) {
+ // cerr << "Removed " << dbg_size0-dbg_size1 << " memo entries.\n";
+ }
+ }
+};
+
+class InputState {
+ // InputState correponds to an input
+ // Input record batches are queued up in InputState until processed and
+ // turned into output record batches.
+
+ public:
+ InputState(const std::shared_ptr<arrow::Schema>& schema,
+ const std::string& time_col_name, const std::string& key_col_name,
Review Comment:
Instead of converting from `FieldRef` to `std::string` you should convert from `FieldRef` to ints. Admittedly, this wouldn't work today because of the shortcut of assuming the key has the same name in each input.
##########
cpp/src/arrow/compute/exec/asof_join_node.cc:
##########
@@ -0,0 +1,800 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+#include <iostream>
+#include <unordered_map>
+
+#include <arrow/api.h>
+#include <arrow/compute/api.h>
+#include <arrow/util/optional.h>
+#include "arrow/compute/exec/asof_join.h"
+#include "arrow/compute/exec/exec_plan.h"
+#include "arrow/compute/exec/options.h"
+#include "arrow/compute/exec/schema_util.h"
+#include "arrow/compute/exec/util.h"
+#include "arrow/result.h"
+#include "arrow/status.h"
+#include "arrow/util/checked_cast.h"
+#include "arrow/util/future.h"
+#include "arrow/util/make_unique.h"
+
+#include <condition_variable>
+#include <mutex>
+#include <thread>
+
+namespace arrow {
+namespace compute {
+
+/**
+ * Simple implementation for an unbound concurrent queue
+ */
+template <class T>
+class ConcurrentQueue {
+ public:
+ T pop() {
+ std::unique_lock<std::mutex> lock(mutex_);
+ cond_.wait(lock, [&] { return !queue_.empty(); });
+ auto item = queue_.front();
+ queue_.pop();
+ return item;
+ }
+
+ void push(const T& item) {
+ std::unique_lock<std::mutex> lock(mutex_);
+ queue_.push(item);
+ cond_.notify_one();
+ }
+
+ util::optional<T> try_pop() {
+ // Try to pop the oldest value from the queue (or return nullopt if none)
+ std::unique_lock<std::mutex> lock(mutex_);
+ if (queue_.empty()) {
+ return util::nullopt;
+ } else {
+ auto item = queue_.front();
+ queue_.pop();
+ return item;
+ }
+ }
+
+ bool empty() const {
+ std::unique_lock<std::mutex> lock(mutex_);
+ return queue_.empty();
+ }
+
+ // Un-synchronized access to front
+ // For this to be "safe":
+ // 1) the caller logically guarantees that queue is not empty
+ // 2) pop/try_pop cannot be called concurrently with this
+ const T& unsync_front() const { return queue_.front(); }
+
+ private:
+ std::queue<T> queue_;
+ mutable std::mutex mutex_;
+ std::condition_variable cond_;
+};
+
+struct MemoStore {
+ // Stores last known values for all the keys
+
+ struct Entry {
+ // Timestamp associated with the entry
+ int64_t _time;
+
+ // Batch associated with the entry (perf is probably OK for this; batches change
+ // rarely)
+ std::shared_ptr<arrow::RecordBatch> _batch;
+
+ // Row associated with the entry
+ row_index_t _row;
+ };
+
+ std::unordered_map<KeyType, Entry> _entries;
+
+ void store(const std::shared_ptr<RecordBatch>& batch, row_index_t row, int64_t time,
+ KeyType key) {
+ auto& e = _entries[key];
+ // that we can do this assignment optionally, is why we
+ // can get array 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;
+ }
+
+ util::optional<const Entry*> get_entry_for_key(KeyType key) const {
+ auto e = _entries.find(key);
+ if (_entries.end() == e) return util::nullopt;
+ return util::optional<const Entry*>(&e->second);
+ }
+
+ void remove_entries_with_lesser_time(int64_t ts) {
+ size_t dbg_size0 = _entries.size();
+ for (auto e = _entries.begin(); e != _entries.end();)
+ if (e->second._time < ts)
+ e = _entries.erase(e);
+ else
+ ++e;
+ size_t dbg_size1 = _entries.size();
+ if (dbg_size1 < dbg_size0) {
+ // cerr << "Removed " << dbg_size0-dbg_size1 << " memo entries.\n";
+ }
+ }
+};
+
+class InputState {
+ // InputState correponds to an input
+ // Input record batches are queued up in InputState until processed and
+ // turned into output record batches.
+
+ public:
+ InputState(const std::shared_ptr<arrow::Schema>& schema,
+ const std::string& time_col_name, const std::string& key_col_name,
+ util::optional<KeyType> wildcard_key)
+ : queue_(),
+ wildcard_key_(wildcard_key),
+ schema_(schema),
+ time_col_index_(
+ schema->GetFieldIndex(time_col_name)), // TODO: handle missing field name
+ key_col_index_(schema->GetFieldIndex(key_col_name)) {}
+
+ col_index_t init_src_to_dst_mapping(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 && is_time_or_key_column(i)))
+ src_to_dst_[i] = dst_offset++;
+ return dst_offset;
+ }
+
+ const util::optional<col_index_t>& map_src_to_dst(col_index_t src) const {
+ return src_to_dst_[src];
+ }
+
+ bool is_time_or_key_column(col_index_t i) const {
+ assert(i < schema_->num_fields());
+ return (i == time_col_index_) || (i == key_col_index_);
+ }
+
+ // Gets the latest row index, assuming the queue isn't empty
+ row_index_t get_latest_row() const { return latest_ref_row_; }
+
+ bool empty() const {
+ if (latest_ref_row_ > 0)
+ return false; // cannot be empty if ref row is >0 -- can avoid slow queue lock
+ // below
+ return queue_.empty();
+ }
+
+ int countbatches_processed_() const { return batches_processed_; }
+ int count_total_batches() const { return total_batches_; }
+
+ // Gets latest batch (precondition: must not be empty)
+ const std::shared_ptr<arrow::RecordBatch>& get_latest_batch() const {
+ return queue_.unsync_front();
+ }
+ KeyType get_latest_key() const {
+ return queue_.unsync_front()
+ ->column_data(key_col_index_)
+ ->GetValues<KeyType>(1)[latest_ref_row_];
+ }
+ int64_t get_latest_time() const {
+ return queue_.unsync_front()
+ ->column_data(time_col_index_)
+ ->GetValues<int64_t>(1)[latest_ref_row_];
+ }
+
+ bool finished() const { return batches_processed_ == total_batches_; }
+
+ bool advance() {
+ // 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();
+ if (have_active_batch) {
+ // If we have an active batch
+ if (++latest_ref_row_ >= (row_index_t)queue_.unsync_front()->num_rows()) {
+ // hit the end of the batch, need to get the next batch if possible.
+ ++batches_processed_;
+ latest_ref_row_ = 0;
+ have_active_batch &= !queue_.try_pop();
+ if (have_active_batch)
+ assert(queue_.unsync_front()->num_rows() > 0); // empty batches disallowed
+ }
+ }
+ return have_active_batch;
+ }
+
+ // Advance the data to be immediately past the specified TS, updating latest and
+ // latest_ref_row to the latest data prior to that immediate just past Returns true if
+ // updates were made, false if not.
+ bool advance_and_memoize(int64_t ts) {
+ // Advance the right side row index until we reach the latest right row (for each key)
+ // for the given left timestamp.
+
+ // Check if already updated for TS (or if there is no latest)
+ if (empty()) return false; // can't advance if empty
+ auto latest_time = get_latest_time();
+ if (latest_time > ts) return false; // already advanced
+
+ // Not updated. Try to update and possibly advance.
+ bool updated = false;
+ do {
+ latest_time = get_latest_time();
+ // 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 <= ts) {
+ memo_.store(get_latest_batch(), latest_ref_row_, latest_time, get_latest_key());
+ } else {
+ break; // hit a future timestamp -- done updating for now
+ }
+ updated = true;
+ } while (advance());
+ return updated;
+ }
+
+ void push(const std::shared_ptr<arrow::RecordBatch>& rb) {
+ if (rb->num_rows() > 0) {
+ queue_.push(rb);
+ } else {
+ ++batches_processed_; // don't enqueue empty batches, just record as processed
+ }
+ }
+
+ util::optional<const MemoStore::Entry*> get_memo_entry_for_key(KeyType key) {
+ auto r = memo_.get_entry_for_key(key);
+ if (r.has_value()) return r;
+ if (wildcard_key_.has_value()) r = memo_.get_entry_for_key(*wildcard_key_);
+ return r;
+ }
+
+ util::optional<int64_t> get_memo_time_for_key(KeyType key) {
+ auto r = get_memo_entry_for_key(key);
+ return r.has_value() ? util::make_optional((*r)->_time) : util::nullopt;
+ }
+
+ void remove_memo_entries_with_lesser_time(int64_t ts) {
+ memo_.remove_entries_with_lesser_time(ts);
+ }
+
+ const std::shared_ptr<Schema>& get_schema() const { return schema_; }
+
+ void set_total_batches(int n) {
+ assert(n >= 0);
+ assert(total_batches_ == -1); // shouldn't be set more than once
+ total_batches_ = n;
+ }
+
+ private:
+ // Pending record batches. The latest is the front. Batches cannot be empty.
+ ConcurrentQueue<std::shared_ptr<RecordBatch>> queue_;
+
+ // Wildcard key for this input, if applicable.
+ util::optional<KeyType> wildcard_key_;
Review Comment:
Can you either provide some more documentation on what this is or remove it until you later make it configurable?
##########
cpp/src/arrow/compute/exec/asof_join_node.cc:
##########
@@ -0,0 +1,800 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+#include <iostream>
+#include <unordered_map>
+
+#include <arrow/api.h>
+#include <arrow/compute/api.h>
+#include <arrow/util/optional.h>
+#include "arrow/compute/exec/asof_join.h"
+#include "arrow/compute/exec/exec_plan.h"
+#include "arrow/compute/exec/options.h"
+#include "arrow/compute/exec/schema_util.h"
+#include "arrow/compute/exec/util.h"
+#include "arrow/result.h"
+#include "arrow/status.h"
+#include "arrow/util/checked_cast.h"
+#include "arrow/util/future.h"
+#include "arrow/util/make_unique.h"
+
+#include <condition_variable>
+#include <mutex>
+#include <thread>
+
+namespace arrow {
+namespace compute {
+
+/**
+ * Simple implementation for an unbound concurrent queue
+ */
+template <class T>
+class ConcurrentQueue {
+ public:
+ T pop() {
+ std::unique_lock<std::mutex> lock(mutex_);
+ cond_.wait(lock, [&] { return !queue_.empty(); });
+ auto item = queue_.front();
+ queue_.pop();
+ return item;
+ }
+
+ void push(const T& item) {
+ std::unique_lock<std::mutex> lock(mutex_);
+ queue_.push(item);
+ cond_.notify_one();
+ }
+
+ util::optional<T> try_pop() {
+ // Try to pop the oldest value from the queue (or return nullopt if none)
+ std::unique_lock<std::mutex> lock(mutex_);
+ if (queue_.empty()) {
+ return util::nullopt;
+ } else {
+ auto item = queue_.front();
+ queue_.pop();
+ return item;
+ }
+ }
+
+ bool empty() const {
+ std::unique_lock<std::mutex> lock(mutex_);
+ return queue_.empty();
+ }
+
+ // Un-synchronized access to front
+ // For this to be "safe":
+ // 1) the caller logically guarantees that queue is not empty
+ // 2) pop/try_pop cannot be called concurrently with this
+ const T& unsync_front() const { return queue_.front(); }
+
+ private:
+ std::queue<T> queue_;
+ mutable std::mutex mutex_;
+ std::condition_variable cond_;
+};
+
+struct MemoStore {
+ // Stores last known values for all the keys
+
+ struct Entry {
+ // Timestamp associated with the entry
+ int64_t _time;
+
+ // Batch associated with the entry (perf is probably OK for this; batches change
+ // rarely)
+ std::shared_ptr<arrow::RecordBatch> _batch;
+
+ // Row associated with the entry
+ row_index_t _row;
+ };
+
+ std::unordered_map<KeyType, Entry> _entries;
+
+ void store(const std::shared_ptr<RecordBatch>& batch, row_index_t row, int64_t time,
+ KeyType key) {
+ auto& e = _entries[key];
+ // that we can do this assignment optionally, is why we
+ // can get array 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;
+ }
+
+ util::optional<const Entry*> get_entry_for_key(KeyType key) const {
+ auto e = _entries.find(key);
+ if (_entries.end() == e) return util::nullopt;
+ return util::optional<const Entry*>(&e->second);
+ }
+
+ void remove_entries_with_lesser_time(int64_t ts) {
+ size_t dbg_size0 = _entries.size();
+ for (auto e = _entries.begin(); e != _entries.end();)
+ if (e->second._time < ts)
+ e = _entries.erase(e);
+ else
+ ++e;
+ size_t dbg_size1 = _entries.size();
+ if (dbg_size1 < dbg_size0) {
+ // cerr << "Removed " << dbg_size0-dbg_size1 << " memo entries.\n";
+ }
+ }
+};
+
+class InputState {
+ // InputState correponds to an input
+ // Input record batches are queued up in InputState until processed and
+ // turned into output record batches.
+
+ public:
+ InputState(const std::shared_ptr<arrow::Schema>& schema,
+ const std::string& time_col_name, const std::string& key_col_name,
+ util::optional<KeyType> wildcard_key)
+ : queue_(),
+ wildcard_key_(wildcard_key),
+ schema_(schema),
+ time_col_index_(
+ schema->GetFieldIndex(time_col_name)), // TODO: handle missing field name
+ key_col_index_(schema->GetFieldIndex(key_col_name)) {}
+
+ col_index_t init_src_to_dst_mapping(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 && is_time_or_key_column(i)))
+ src_to_dst_[i] = dst_offset++;
+ return dst_offset;
+ }
+
+ const util::optional<col_index_t>& map_src_to_dst(col_index_t src) const {
+ return src_to_dst_[src];
+ }
+
+ bool is_time_or_key_column(col_index_t i) const {
+ assert(i < schema_->num_fields());
+ return (i == time_col_index_) || (i == key_col_index_);
+ }
+
+ // Gets the latest row index, assuming the queue isn't empty
+ row_index_t get_latest_row() const { return latest_ref_row_; }
+
+ bool empty() const {
+ if (latest_ref_row_ > 0)
+ return false; // cannot be empty if ref row is >0 -- can avoid slow queue lock
+ // below
+ return queue_.empty();
+ }
+
+ int countbatches_processed_() const { return batches_processed_; }
+ int count_total_batches() const { return total_batches_; }
+
+ // Gets latest batch (precondition: must not be empty)
+ const std::shared_ptr<arrow::RecordBatch>& get_latest_batch() const {
+ return queue_.unsync_front();
+ }
+ KeyType get_latest_key() const {
+ return queue_.unsync_front()
+ ->column_data(key_col_index_)
+ ->GetValues<KeyType>(1)[latest_ref_row_];
+ }
+ int64_t get_latest_time() const {
+ return queue_.unsync_front()
+ ->column_data(time_col_index_)
+ ->GetValues<int64_t>(1)[latest_ref_row_];
+ }
+
+ bool finished() const { return batches_processed_ == total_batches_; }
+
+ bool advance() {
+ // 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();
+ if (have_active_batch) {
+ // If we have an active batch
+ if (++latest_ref_row_ >= (row_index_t)queue_.unsync_front()->num_rows()) {
+ // hit the end of the batch, need to get the next batch if possible.
+ ++batches_processed_;
+ latest_ref_row_ = 0;
+ have_active_batch &= !queue_.try_pop();
+ if (have_active_batch)
+ assert(queue_.unsync_front()->num_rows() > 0); // empty batches disallowed
+ }
+ }
+ return have_active_batch;
+ }
+
+ // Advance the data to be immediately past the specified TS, updating latest and
+ // latest_ref_row to the latest data prior to that immediate just past Returns true if
+ // updates were made, false if not.
+ bool advance_and_memoize(int64_t ts) {
+ // Advance the right side row index until we reach the latest right row (for each key)
+ // for the given left timestamp.
+
+ // Check if already updated for TS (or if there is no latest)
+ if (empty()) return false; // can't advance if empty
+ auto latest_time = get_latest_time();
+ if (latest_time > ts) return false; // already advanced
+
+ // Not updated. Try to update and possibly advance.
+ bool updated = false;
+ do {
+ latest_time = get_latest_time();
+ // 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 <= ts) {
+ memo_.store(get_latest_batch(), latest_ref_row_, latest_time, get_latest_key());
+ } else {
+ break; // hit a future timestamp -- done updating for now
+ }
+ updated = true;
+ } while (advance());
+ return updated;
+ }
+
+ void push(const std::shared_ptr<arrow::RecordBatch>& rb) {
+ if (rb->num_rows() > 0) {
+ queue_.push(rb);
+ } else {
+ ++batches_processed_; // don't enqueue empty batches, just record as processed
+ }
+ }
+
+ util::optional<const MemoStore::Entry*> get_memo_entry_for_key(KeyType key) {
+ auto r = memo_.get_entry_for_key(key);
+ if (r.has_value()) return r;
+ if (wildcard_key_.has_value()) r = memo_.get_entry_for_key(*wildcard_key_);
+ return r;
+ }
+
+ util::optional<int64_t> get_memo_time_for_key(KeyType key) {
+ auto r = get_memo_entry_for_key(key);
+ return r.has_value() ? util::make_optional((*r)->_time) : util::nullopt;
+ }
+
+ void remove_memo_entries_with_lesser_time(int64_t ts) {
+ memo_.remove_entries_with_lesser_time(ts);
+ }
+
+ const std::shared_ptr<Schema>& get_schema() const { return schema_; }
+
+ void set_total_batches(int n) {
+ assert(n >= 0);
+ assert(total_batches_ == -1); // shouldn't be set more than once
+ total_batches_ = n;
+ }
+
+ private:
+ // Pending record batches. The latest is the front. Batches cannot be empty.
+ ConcurrentQueue<std::shared_ptr<RecordBatch>> queue_;
+
+ // Wildcard key for this input, if applicable.
+ util::optional<KeyType> wildcard_key_;
+
+ // Schema associated with the input
+ std::shared_ptr<Schema> schema_;
+
+ // Total number of batches (only int because InputFinished uses int)
+ int total_batches_ = -1;
+
+ // Number of batches processed so far (only int because InputFinished uses int)
+ int batches_processed_ = 0;
+
+ // Index of the time col
+ col_index_t time_col_index_;
+
+ // Index of the key col
+ col_index_t key_col_index_;
+
+ // Index of the latest row reference within; if >0 then queue_ cannot be empty
+ row_index_t latest_ref_row_ =
+ 0; // must be < queue_.front()->num_rows() if queue_ is non-empty
+
+ // Stores latest known values for the various keys
+ MemoStore memo_;
+
+ // Mapping of source columns to destination columns
+ std::vector<util::optional<col_index_t>> src_to_dst_;
+};
+
+template <size_t MAX_TABLES>
+struct CompositeReferenceRow {
+ struct Entry {
+ arrow::RecordBatch* batch; // can be NULL if there's no value
+ row_index_t row;
+ };
+ Entry refs[MAX_TABLES];
+};
+
+// A table of composite reference rows. Rows maintain pointers to the
+// constituent record batches, but the overall table retains shared_ptr
+// references to ensure memory remains resident while the table is live.
+//
+// The main reason for this is that, especially for wide tables, joins
+// are effectively row-oriented, rather than column-oriented. Separating
+// the join part from the columnar materialization part simplifies the
+// logic around data types and increases efficiency.
+//
+// We don't put the shared_ptr's into the rows for efficiency reasons.
+template <size_t MAX_TABLES>
+class CompositeReferenceTable {
+ public:
+ explicit CompositeReferenceTable(size_t n_tables) : n_tables_(n_tables) {
+ assert(n_tables_ >= 1);
+ assert(n_tables_ <= MAX_TABLES);
+ }
+
+ size_t n_rows() const { return rows_.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<std::unique_ptr<InputState>>& in, int64_t tolerance) {
+ assert(in.size() == n_tables_);
+
+ // Get the LHS key
+ KeyType key = in[0]->get_latest_key();
+
+ // Add row and setup LHS
+ // (the LHS state comes just from the latest row of the LHS table)
+ assert(!in[0]->empty());
+ const std::shared_ptr<arrow::RecordBatch>& lhs_latest_batch =
+ in[0]->get_latest_batch();
+ row_index_t lhs_latest_row = in[0]->get_latest_row();
+ int64_t lhs_latest_time = in[0]->get_latest_time();
+ 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 = rows_.size() + new_batch_size;
+ if (rows_.capacity() < new_capacity) rows_.reserve(new_capacity);
Review Comment:
Can this ever be false? Not a problem, just curious.
##########
cpp/src/arrow/compute/exec/asof_join_node.cc:
##########
@@ -0,0 +1,800 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+#include <iostream>
+#include <unordered_map>
+
+#include <arrow/api.h>
+#include <arrow/compute/api.h>
+#include <arrow/util/optional.h>
+#include "arrow/compute/exec/asof_join.h"
+#include "arrow/compute/exec/exec_plan.h"
+#include "arrow/compute/exec/options.h"
+#include "arrow/compute/exec/schema_util.h"
+#include "arrow/compute/exec/util.h"
+#include "arrow/result.h"
+#include "arrow/status.h"
+#include "arrow/util/checked_cast.h"
+#include "arrow/util/future.h"
+#include "arrow/util/make_unique.h"
+
+#include <condition_variable>
+#include <mutex>
+#include <thread>
+
+namespace arrow {
+namespace compute {
+
+/**
+ * Simple implementation for an unbound concurrent queue
+ */
+template <class T>
+class ConcurrentQueue {
+ public:
+ T pop() {
+ std::unique_lock<std::mutex> lock(mutex_);
+ cond_.wait(lock, [&] { return !queue_.empty(); });
+ auto item = queue_.front();
+ queue_.pop();
+ return item;
+ }
+
+ void push(const T& item) {
+ std::unique_lock<std::mutex> lock(mutex_);
+ queue_.push(item);
+ cond_.notify_one();
+ }
+
+ util::optional<T> try_pop() {
+ // Try to pop the oldest value from the queue (or return nullopt if none)
+ std::unique_lock<std::mutex> lock(mutex_);
+ if (queue_.empty()) {
+ return util::nullopt;
+ } else {
+ auto item = queue_.front();
+ queue_.pop();
+ return item;
+ }
+ }
+
+ bool empty() const {
+ std::unique_lock<std::mutex> lock(mutex_);
+ return queue_.empty();
+ }
+
+ // Un-synchronized access to front
+ // For this to be "safe":
+ // 1) the caller logically guarantees that queue is not empty
+ // 2) pop/try_pop cannot be called concurrently with this
+ const T& unsync_front() const { return queue_.front(); }
+
+ private:
+ std::queue<T> queue_;
+ mutable std::mutex mutex_;
+ std::condition_variable cond_;
+};
+
+struct MemoStore {
+ // Stores last known values for all the keys
+
+ struct Entry {
+ // Timestamp associated with the entry
+ int64_t _time;
+
+ // Batch associated with the entry (perf is probably OK for this; batches change
+ // rarely)
+ std::shared_ptr<arrow::RecordBatch> _batch;
+
+ // Row associated with the entry
+ row_index_t _row;
+ };
+
+ std::unordered_map<KeyType, Entry> _entries;
+
+ void store(const std::shared_ptr<RecordBatch>& batch, row_index_t row, int64_t time,
+ KeyType key) {
+ auto& e = _entries[key];
+ // that we can do this assignment optionally, is why we
+ // can get array 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;
+ }
+
+ util::optional<const Entry*> get_entry_for_key(KeyType key) const {
+ auto e = _entries.find(key);
+ if (_entries.end() == e) return util::nullopt;
+ return util::optional<const Entry*>(&e->second);
+ }
+
+ void remove_entries_with_lesser_time(int64_t ts) {
+ size_t dbg_size0 = _entries.size();
+ for (auto e = _entries.begin(); e != _entries.end();)
+ if (e->second._time < ts)
+ e = _entries.erase(e);
+ else
+ ++e;
+ size_t dbg_size1 = _entries.size();
+ if (dbg_size1 < dbg_size0) {
+ // cerr << "Removed " << dbg_size0-dbg_size1 << " memo entries.\n";
+ }
+ }
+};
+
+class InputState {
+ // InputState correponds to an input
+ // Input record batches are queued up in InputState until processed and
+ // turned into output record batches.
+
+ public:
+ InputState(const std::shared_ptr<arrow::Schema>& schema,
+ const std::string& time_col_name, const std::string& key_col_name,
+ util::optional<KeyType> wildcard_key)
+ : queue_(),
+ wildcard_key_(wildcard_key),
+ schema_(schema),
+ time_col_index_(
+ schema->GetFieldIndex(time_col_name)), // TODO: handle missing field name
+ key_col_index_(schema->GetFieldIndex(key_col_name)) {}
+
+ col_index_t init_src_to_dst_mapping(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 && is_time_or_key_column(i)))
+ src_to_dst_[i] = dst_offset++;
+ return dst_offset;
+ }
+
+ const util::optional<col_index_t>& map_src_to_dst(col_index_t src) const {
+ return src_to_dst_[src];
+ }
+
+ bool is_time_or_key_column(col_index_t i) const {
+ assert(i < schema_->num_fields());
+ return (i == time_col_index_) || (i == key_col_index_);
+ }
+
+ // Gets the latest row index, assuming the queue isn't empty
+ row_index_t get_latest_row() const { return latest_ref_row_; }
+
+ bool empty() const {
+ if (latest_ref_row_ > 0)
+ return false; // cannot be empty if ref row is >0 -- can avoid slow queue lock
+ // below
+ return queue_.empty();
+ }
+
+ int countbatches_processed_() const { return batches_processed_; }
+ int count_total_batches() const { return total_batches_; }
+
+ // Gets latest batch (precondition: must not be empty)
+ const std::shared_ptr<arrow::RecordBatch>& get_latest_batch() const {
+ return queue_.unsync_front();
+ }
+ KeyType get_latest_key() const {
+ return queue_.unsync_front()
+ ->column_data(key_col_index_)
+ ->GetValues<KeyType>(1)[latest_ref_row_];
+ }
+ int64_t get_latest_time() const {
+ return queue_.unsync_front()
+ ->column_data(time_col_index_)
+ ->GetValues<int64_t>(1)[latest_ref_row_];
+ }
+
+ bool finished() const { return batches_processed_ == total_batches_; }
+
+ bool advance() {
+ // 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();
+ if (have_active_batch) {
+ // If we have an active batch
+ if (++latest_ref_row_ >= (row_index_t)queue_.unsync_front()->num_rows()) {
+ // hit the end of the batch, need to get the next batch if possible.
+ ++batches_processed_;
+ latest_ref_row_ = 0;
+ have_active_batch &= !queue_.try_pop();
+ if (have_active_batch)
+ assert(queue_.unsync_front()->num_rows() > 0); // empty batches disallowed
+ }
+ }
+ return have_active_batch;
+ }
+
+ // Advance the data to be immediately past the specified TS, updating latest and
+ // latest_ref_row to the latest data prior to that immediate just past Returns true if
+ // updates were made, false if not.
+ bool advance_and_memoize(int64_t ts) {
+ // Advance the right side row index until we reach the latest right row (for each key)
+ // for the given left timestamp.
+
+ // Check if already updated for TS (or if there is no latest)
+ if (empty()) return false; // can't advance if empty
+ auto latest_time = get_latest_time();
+ if (latest_time > ts) return false; // already advanced
+
+ // Not updated. Try to update and possibly advance.
+ bool updated = false;
+ do {
+ latest_time = get_latest_time();
+ // 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 <= ts) {
+ memo_.store(get_latest_batch(), latest_ref_row_, latest_time, get_latest_key());
+ } else {
+ break; // hit a future timestamp -- done updating for now
+ }
+ updated = true;
+ } while (advance());
+ return updated;
+ }
+
+ void push(const std::shared_ptr<arrow::RecordBatch>& rb) {
+ if (rb->num_rows() > 0) {
+ queue_.push(rb);
+ } else {
+ ++batches_processed_; // don't enqueue empty batches, just record as processed
+ }
+ }
+
+ util::optional<const MemoStore::Entry*> get_memo_entry_for_key(KeyType key) {
+ auto r = memo_.get_entry_for_key(key);
+ if (r.has_value()) return r;
+ if (wildcard_key_.has_value()) r = memo_.get_entry_for_key(*wildcard_key_);
+ return r;
+ }
+
+ util::optional<int64_t> get_memo_time_for_key(KeyType key) {
+ auto r = get_memo_entry_for_key(key);
+ return r.has_value() ? util::make_optional((*r)->_time) : util::nullopt;
+ }
+
+ void remove_memo_entries_with_lesser_time(int64_t ts) {
+ memo_.remove_entries_with_lesser_time(ts);
+ }
+
+ const std::shared_ptr<Schema>& get_schema() const { return schema_; }
+
+ void set_total_batches(int n) {
+ assert(n >= 0);
+ assert(total_batches_ == -1); // shouldn't be set more than once
+ total_batches_ = n;
+ }
+
+ private:
+ // Pending record batches. The latest is the front. Batches cannot be empty.
+ ConcurrentQueue<std::shared_ptr<RecordBatch>> queue_;
+
+ // Wildcard key for this input, if applicable.
+ util::optional<KeyType> wildcard_key_;
+
+ // Schema associated with the input
+ std::shared_ptr<Schema> schema_;
+
+ // Total number of batches (only int because InputFinished uses int)
+ int total_batches_ = -1;
+
+ // Number of batches processed so far (only int because InputFinished uses int)
+ int batches_processed_ = 0;
+
+ // Index of the time col
+ col_index_t time_col_index_;
+
+ // Index of the key col
+ col_index_t key_col_index_;
+
+ // Index of the latest row reference within; if >0 then queue_ cannot be empty
+ row_index_t latest_ref_row_ =
+ 0; // must be < queue_.front()->num_rows() if queue_ is non-empty
+
+ // Stores latest known values for the various keys
+ MemoStore memo_;
+
+ // Mapping of source columns to destination columns
+ std::vector<util::optional<col_index_t>> src_to_dst_;
+};
+
+template <size_t MAX_TABLES>
+struct CompositeReferenceRow {
+ struct Entry {
+ arrow::RecordBatch* batch; // can be NULL if there's no value
+ row_index_t row;
+ };
+ Entry refs[MAX_TABLES];
+};
+
+// A table of composite reference rows. Rows maintain pointers to the
+// constituent record batches, but the overall table retains shared_ptr
+// references to ensure memory remains resident while the table is live.
+//
+// The main reason for this is that, especially for wide tables, joins
+// are effectively row-oriented, rather than column-oriented. Separating
+// the join part from the columnar materialization part simplifies the
+// logic around data types and increases efficiency.
+//
+// We don't put the shared_ptr's into the rows for efficiency reasons.
+template <size_t MAX_TABLES>
+class CompositeReferenceTable {
+ public:
+ explicit CompositeReferenceTable(size_t n_tables) : n_tables_(n_tables) {
+ assert(n_tables_ >= 1);
+ assert(n_tables_ <= MAX_TABLES);
+ }
+
+ size_t n_rows() const { return rows_.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<std::unique_ptr<InputState>>& in, int64_t tolerance) {
+ assert(in.size() == n_tables_);
+
+ // Get the LHS key
+ KeyType key = in[0]->get_latest_key();
+
+ // Add row and setup LHS
+ // (the LHS state comes just from the latest row of the LHS table)
+ assert(!in[0]->empty());
+ const std::shared_ptr<arrow::RecordBatch>& lhs_latest_batch =
+ in[0]->get_latest_batch();
+ row_index_t lhs_latest_row = in[0]->get_latest_row();
+ int64_t lhs_latest_time = in[0]->get_latest_time();
+ 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 = rows_.size() + new_batch_size;
+ if (rows_.capacity() < new_capacity) rows_.reserve(new_capacity);
+ }
+ rows_.resize(rows_.size() + 1);
+ auto& row = rows_.back();
+ row.refs[0].batch = lhs_latest_batch.get();
+ row.refs[0].row = lhs_latest_row;
+ add_record_batch_ref(lhs_latest_batch);
+
+ // 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) {
+ util::optional<const MemoStore::Entry*> opt_entry =
+ in[i]->get_memo_entry_for_key(key);
+ if (opt_entry.has_value()) {
+ assert(*opt_entry);
Review Comment:
```suggestion
DCHECK(*opt_entry);
```
##########
cpp/src/arrow/compute/exec/asof_join_node.cc:
##########
@@ -0,0 +1,800 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+#include <iostream>
+#include <unordered_map>
+
+#include <arrow/api.h>
+#include <arrow/compute/api.h>
+#include <arrow/util/optional.h>
+#include "arrow/compute/exec/asof_join.h"
+#include "arrow/compute/exec/exec_plan.h"
+#include "arrow/compute/exec/options.h"
+#include "arrow/compute/exec/schema_util.h"
+#include "arrow/compute/exec/util.h"
+#include "arrow/result.h"
+#include "arrow/status.h"
+#include "arrow/util/checked_cast.h"
+#include "arrow/util/future.h"
+#include "arrow/util/make_unique.h"
+
+#include <condition_variable>
+#include <mutex>
+#include <thread>
+
+namespace arrow {
+namespace compute {
+
+/**
+ * Simple implementation for an unbound concurrent queue
+ */
+template <class T>
+class ConcurrentQueue {
+ public:
+ T pop() {
+ std::unique_lock<std::mutex> lock(mutex_);
+ cond_.wait(lock, [&] { return !queue_.empty(); });
+ auto item = queue_.front();
+ queue_.pop();
+ return item;
+ }
+
+ void push(const T& item) {
+ std::unique_lock<std::mutex> lock(mutex_);
+ queue_.push(item);
+ cond_.notify_one();
+ }
+
+ util::optional<T> try_pop() {
+ // Try to pop the oldest value from the queue (or return nullopt if none)
+ std::unique_lock<std::mutex> lock(mutex_);
+ if (queue_.empty()) {
+ return util::nullopt;
+ } else {
+ auto item = queue_.front();
+ queue_.pop();
+ return item;
+ }
+ }
+
+ bool empty() const {
+ std::unique_lock<std::mutex> lock(mutex_);
+ return queue_.empty();
+ }
+
+ // Un-synchronized access to front
+ // For this to be "safe":
+ // 1) the caller logically guarantees that queue is not empty
+ // 2) pop/try_pop cannot be called concurrently with this
+ const T& unsync_front() const { return queue_.front(); }
+
+ private:
+ std::queue<T> queue_;
+ mutable std::mutex mutex_;
+ std::condition_variable cond_;
+};
+
+struct MemoStore {
+ // Stores last known values for all the keys
+
+ struct Entry {
+ // Timestamp associated with the entry
+ int64_t _time;
+
+ // Batch associated with the entry (perf is probably OK for this; batches change
+ // rarely)
+ std::shared_ptr<arrow::RecordBatch> _batch;
+
+ // Row associated with the entry
+ row_index_t _row;
+ };
+
+ std::unordered_map<KeyType, Entry> _entries;
+
+ void store(const std::shared_ptr<RecordBatch>& batch, row_index_t row, int64_t time,
Review Comment:
```suggestion
void Store(const std::shared_ptr<RecordBatch>& batch, row_index_t row, int64_t time,
```
##########
cpp/src/arrow/compute/exec/options.h:
##########
@@ -361,6 +361,19 @@ class ARROW_EXPORT HashJoinNodeOptions : public ExecNodeOptions {
Expression filter;
};
+class ARROW_EXPORT AsofJoinNodeOptions : public ExecNodeOptions {
+ public:
+ AsofJoinNodeOptions(FieldRef time, FieldRef keys, int64_t tolerance)
+ : time(std::move(time)), keys(std::move(keys)), tolerance(tolerance) {}
+
+ // time column
+ FieldRef time;
Review Comment:
Is `time` the best name here? While I agree that this join is typically used for time series analysis I believe the only requirement is that this be an ordered column present in all inputs. Maybe `asof_key` or something like that?.
##########
cpp/src/arrow/compute/exec/asof_join_node.cc:
##########
@@ -0,0 +1,800 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+#include <iostream>
+#include <unordered_map>
+
+#include <arrow/api.h>
+#include <arrow/compute/api.h>
+#include <arrow/util/optional.h>
+#include "arrow/compute/exec/asof_join.h"
+#include "arrow/compute/exec/exec_plan.h"
+#include "arrow/compute/exec/options.h"
+#include "arrow/compute/exec/schema_util.h"
+#include "arrow/compute/exec/util.h"
+#include "arrow/result.h"
+#include "arrow/status.h"
+#include "arrow/util/checked_cast.h"
+#include "arrow/util/future.h"
+#include "arrow/util/make_unique.h"
+
+#include <condition_variable>
+#include <mutex>
+#include <thread>
+
+namespace arrow {
+namespace compute {
+
+/**
+ * Simple implementation for an unbound concurrent queue
+ */
+template <class T>
+class ConcurrentQueue {
+ public:
+ T pop() {
+ std::unique_lock<std::mutex> lock(mutex_);
+ cond_.wait(lock, [&] { return !queue_.empty(); });
+ auto item = queue_.front();
+ queue_.pop();
+ return item;
+ }
+
+ void push(const T& item) {
Review Comment:
```suggestion
void Push(const T& item) {
```
##########
cpp/src/arrow/compute/exec/asof_join.h:
##########
@@ -0,0 +1,42 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+#pragma once
+
+#include <arrow/compute/api.h>
+#include <arrow/compute/exec/exec_plan.h>
+#include "arrow/compute/exec/options.h"
+#include "arrow/compute/exec/schema_util.h"
+
+namespace arrow {
+namespace compute {
+
+typedef int32_t KeyType;
Review Comment:
a typedef in a header file is ideally avoided or namespaced inside a struct/class of some kind. This creates `arrow::compute::KeyType` but `KeyType` doesn't make sense across the entire `arrow::compute` namespace.
Admittedly, this is a moot point if we move the contents here into a .cc file.
##########
cpp/src/arrow/compute/exec/asof_join_node.cc:
##########
@@ -0,0 +1,800 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+#include <iostream>
+#include <unordered_map>
+
+#include <arrow/api.h>
+#include <arrow/compute/api.h>
+#include <arrow/util/optional.h>
+#include "arrow/compute/exec/asof_join.h"
+#include "arrow/compute/exec/exec_plan.h"
+#include "arrow/compute/exec/options.h"
+#include "arrow/compute/exec/schema_util.h"
+#include "arrow/compute/exec/util.h"
+#include "arrow/result.h"
+#include "arrow/status.h"
+#include "arrow/util/checked_cast.h"
+#include "arrow/util/future.h"
+#include "arrow/util/make_unique.h"
+
+#include <condition_variable>
+#include <mutex>
+#include <thread>
+
+namespace arrow {
+namespace compute {
+
+/**
+ * Simple implementation for an unbound concurrent queue
+ */
+template <class T>
+class ConcurrentQueue {
+ public:
+ T pop() {
+ std::unique_lock<std::mutex> lock(mutex_);
+ cond_.wait(lock, [&] { return !queue_.empty(); });
+ auto item = queue_.front();
+ queue_.pop();
+ return item;
+ }
+
+ void push(const T& item) {
+ std::unique_lock<std::mutex> lock(mutex_);
+ queue_.push(item);
+ cond_.notify_one();
+ }
+
+ util::optional<T> try_pop() {
+ // Try to pop the oldest value from the queue (or return nullopt if none)
+ std::unique_lock<std::mutex> lock(mutex_);
+ if (queue_.empty()) {
+ return util::nullopt;
+ } else {
+ auto item = queue_.front();
+ queue_.pop();
+ return item;
+ }
+ }
+
+ bool empty() const {
Review Comment:
```suggestion
bool Empty() const {
```
##########
cpp/src/arrow/compute/exec/asof_join_node.cc:
##########
@@ -0,0 +1,800 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+#include <iostream>
+#include <unordered_map>
+
+#include <arrow/api.h>
+#include <arrow/compute/api.h>
+#include <arrow/util/optional.h>
+#include "arrow/compute/exec/asof_join.h"
+#include "arrow/compute/exec/exec_plan.h"
+#include "arrow/compute/exec/options.h"
+#include "arrow/compute/exec/schema_util.h"
+#include "arrow/compute/exec/util.h"
+#include "arrow/result.h"
+#include "arrow/status.h"
+#include "arrow/util/checked_cast.h"
+#include "arrow/util/future.h"
+#include "arrow/util/make_unique.h"
+
+#include <condition_variable>
+#include <mutex>
+#include <thread>
+
+namespace arrow {
+namespace compute {
+
+/**
+ * Simple implementation for an unbound concurrent queue
+ */
+template <class T>
+class ConcurrentQueue {
+ public:
+ T pop() {
+ std::unique_lock<std::mutex> lock(mutex_);
+ cond_.wait(lock, [&] { return !queue_.empty(); });
+ auto item = queue_.front();
+ queue_.pop();
+ return item;
+ }
+
+ void push(const T& item) {
+ std::unique_lock<std::mutex> lock(mutex_);
+ queue_.push(item);
+ cond_.notify_one();
+ }
+
+ util::optional<T> try_pop() {
Review Comment:
```suggestion
util::optional<T> TryPop() {
```
##########
cpp/src/arrow/compute/exec/asof_join_node.cc:
##########
@@ -0,0 +1,800 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+#include <iostream>
+#include <unordered_map>
+
+#include <arrow/api.h>
+#include <arrow/compute/api.h>
+#include <arrow/util/optional.h>
+#include "arrow/compute/exec/asof_join.h"
+#include "arrow/compute/exec/exec_plan.h"
+#include "arrow/compute/exec/options.h"
+#include "arrow/compute/exec/schema_util.h"
+#include "arrow/compute/exec/util.h"
+#include "arrow/result.h"
+#include "arrow/status.h"
+#include "arrow/util/checked_cast.h"
+#include "arrow/util/future.h"
+#include "arrow/util/make_unique.h"
+
+#include <condition_variable>
+#include <mutex>
+#include <thread>
+
+namespace arrow {
+namespace compute {
+
+/**
+ * Simple implementation for an unbound concurrent queue
+ */
+template <class T>
+class ConcurrentQueue {
+ public:
+ T pop() {
+ std::unique_lock<std::mutex> lock(mutex_);
+ cond_.wait(lock, [&] { return !queue_.empty(); });
+ auto item = queue_.front();
+ queue_.pop();
+ return item;
+ }
+
+ void push(const T& item) {
+ std::unique_lock<std::mutex> lock(mutex_);
+ queue_.push(item);
+ cond_.notify_one();
+ }
+
+ util::optional<T> try_pop() {
+ // Try to pop the oldest value from the queue (or return nullopt if none)
+ std::unique_lock<std::mutex> lock(mutex_);
+ if (queue_.empty()) {
+ return util::nullopt;
+ } else {
+ auto item = queue_.front();
+ queue_.pop();
+ return item;
+ }
+ }
+
+ bool empty() const {
+ std::unique_lock<std::mutex> lock(mutex_);
+ return queue_.empty();
+ }
+
+ // Un-synchronized access to front
+ // For this to be "safe":
+ // 1) the caller logically guarantees that queue is not empty
+ // 2) pop/try_pop cannot be called concurrently with this
+ const T& unsync_front() const { return queue_.front(); }
+
+ private:
+ std::queue<T> queue_;
+ mutable std::mutex mutex_;
+ std::condition_variable cond_;
+};
+
+struct MemoStore {
+ // Stores last known values for all the keys
+
+ struct Entry {
+ // Timestamp associated with the entry
+ int64_t _time;
+
+ // Batch associated with the entry (perf is probably OK for this; batches change
+ // rarely)
+ std::shared_ptr<arrow::RecordBatch> _batch;
+
+ // Row associated with the entry
+ row_index_t _row;
+ };
+
+ std::unordered_map<KeyType, Entry> _entries;
+
+ void store(const std::shared_ptr<RecordBatch>& batch, row_index_t row, int64_t time,
+ KeyType key) {
+ auto& e = _entries[key];
+ // that we can do this assignment optionally, is why we
+ // can get array 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;
+ }
+
+ util::optional<const Entry*> get_entry_for_key(KeyType key) const {
+ auto e = _entries.find(key);
+ if (_entries.end() == e) return util::nullopt;
+ return util::optional<const Entry*>(&e->second);
+ }
+
+ void remove_entries_with_lesser_time(int64_t ts) {
Review Comment:
```suggestion
void RemoveEntriesWithLesserTime(int64_t ts) {
```
For brevity I'm going to stop suggesting snake_case -> CamelCase at this point but please fix all the other methods. We try to only use `snake_case` for simple accessors.
##########
cpp/src/arrow/compute/exec/asof_join_node.cc:
##########
@@ -0,0 +1,800 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+#include <iostream>
+#include <unordered_map>
+
+#include <arrow/api.h>
+#include <arrow/compute/api.h>
+#include <arrow/util/optional.h>
+#include "arrow/compute/exec/asof_join.h"
+#include "arrow/compute/exec/exec_plan.h"
+#include "arrow/compute/exec/options.h"
+#include "arrow/compute/exec/schema_util.h"
+#include "arrow/compute/exec/util.h"
+#include "arrow/result.h"
+#include "arrow/status.h"
+#include "arrow/util/checked_cast.h"
+#include "arrow/util/future.h"
+#include "arrow/util/make_unique.h"
+
+#include <condition_variable>
+#include <mutex>
+#include <thread>
+
+namespace arrow {
+namespace compute {
+
+/**
+ * Simple implementation for an unbound concurrent queue
+ */
+template <class T>
+class ConcurrentQueue {
+ public:
+ T pop() {
+ std::unique_lock<std::mutex> lock(mutex_);
+ cond_.wait(lock, [&] { return !queue_.empty(); });
+ auto item = queue_.front();
+ queue_.pop();
+ return item;
+ }
+
+ void push(const T& item) {
+ std::unique_lock<std::mutex> lock(mutex_);
+ queue_.push(item);
+ cond_.notify_one();
+ }
+
+ util::optional<T> try_pop() {
+ // Try to pop the oldest value from the queue (or return nullopt if none)
+ std::unique_lock<std::mutex> lock(mutex_);
+ if (queue_.empty()) {
+ return util::nullopt;
+ } else {
+ auto item = queue_.front();
+ queue_.pop();
+ return item;
+ }
+ }
+
+ bool empty() const {
+ std::unique_lock<std::mutex> lock(mutex_);
+ return queue_.empty();
+ }
+
+ // Un-synchronized access to front
+ // For this to be "safe":
+ // 1) the caller logically guarantees that queue is not empty
+ // 2) pop/try_pop cannot be called concurrently with this
+ const T& unsync_front() const { return queue_.front(); }
+
+ private:
+ std::queue<T> queue_;
+ mutable std::mutex mutex_;
+ std::condition_variable cond_;
+};
+
+struct MemoStore {
+ // Stores last known values for all the keys
+
+ struct Entry {
+ // Timestamp associated with the entry
+ int64_t _time;
+
+ // Batch associated with the entry (perf is probably OK for this; batches change
+ // rarely)
+ std::shared_ptr<arrow::RecordBatch> _batch;
+
+ // Row associated with the entry
+ row_index_t _row;
+ };
+
+ std::unordered_map<KeyType, Entry> _entries;
+
+ void store(const std::shared_ptr<RecordBatch>& batch, row_index_t row, int64_t time,
+ KeyType key) {
+ auto& e = _entries[key];
+ // that we can do this assignment optionally, is why we
+ // can get array 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;
+ }
+
+ util::optional<const Entry*> get_entry_for_key(KeyType key) const {
+ auto e = _entries.find(key);
+ if (_entries.end() == e) return util::nullopt;
+ return util::optional<const Entry*>(&e->second);
+ }
+
+ void remove_entries_with_lesser_time(int64_t ts) {
+ size_t dbg_size0 = _entries.size();
+ for (auto e = _entries.begin(); e != _entries.end();)
+ if (e->second._time < ts)
+ e = _entries.erase(e);
+ else
+ ++e;
+ size_t dbg_size1 = _entries.size();
+ if (dbg_size1 < dbg_size0) {
+ // cerr << "Removed " << dbg_size0-dbg_size1 << " memo entries.\n";
+ }
+ }
+};
+
+class InputState {
+ // InputState correponds to an input
+ // Input record batches are queued up in InputState until processed and
+ // turned into output record batches.
+
+ public:
+ InputState(const std::shared_ptr<arrow::Schema>& schema,
+ const std::string& time_col_name, const std::string& key_col_name,
+ util::optional<KeyType> wildcard_key)
+ : queue_(),
+ wildcard_key_(wildcard_key),
+ schema_(schema),
+ time_col_index_(
+ schema->GetFieldIndex(time_col_name)), // TODO: handle missing field name
+ key_col_index_(schema->GetFieldIndex(key_col_name)) {}
+
+ col_index_t init_src_to_dst_mapping(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 && is_time_or_key_column(i)))
+ src_to_dst_[i] = dst_offset++;
+ return dst_offset;
+ }
+
+ const util::optional<col_index_t>& map_src_to_dst(col_index_t src) const {
+ return src_to_dst_[src];
+ }
+
+ bool is_time_or_key_column(col_index_t i) const {
+ assert(i < schema_->num_fields());
+ return (i == time_col_index_) || (i == key_col_index_);
+ }
+
+ // Gets the latest row index, assuming the queue isn't empty
+ row_index_t get_latest_row() const { return latest_ref_row_; }
+
+ bool empty() const {
+ if (latest_ref_row_ > 0)
+ return false; // cannot be empty if ref row is >0 -- can avoid slow queue lock
+ // below
+ return queue_.empty();
+ }
+
+ int countbatches_processed_() const { return batches_processed_; }
+ int count_total_batches() const { return total_batches_; }
+
+ // Gets latest batch (precondition: must not be empty)
+ const std::shared_ptr<arrow::RecordBatch>& get_latest_batch() const {
+ return queue_.unsync_front();
+ }
+ KeyType get_latest_key() const {
+ return queue_.unsync_front()
+ ->column_data(key_col_index_)
+ ->GetValues<KeyType>(1)[latest_ref_row_];
+ }
+ int64_t get_latest_time() const {
+ return queue_.unsync_front()
+ ->column_data(time_col_index_)
+ ->GetValues<int64_t>(1)[latest_ref_row_];
+ }
+
+ bool finished() const { return batches_processed_ == total_batches_; }
+
+ bool advance() {
+ // 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();
+ if (have_active_batch) {
+ // If we have an active batch
+ if (++latest_ref_row_ >= (row_index_t)queue_.unsync_front()->num_rows()) {
+ // hit the end of the batch, need to get the next batch if possible.
+ ++batches_processed_;
+ latest_ref_row_ = 0;
+ have_active_batch &= !queue_.try_pop();
+ if (have_active_batch)
+ assert(queue_.unsync_front()->num_rows() > 0); // empty batches disallowed
+ }
+ }
+ return have_active_batch;
+ }
+
+ // Advance the data to be immediately past the specified TS, updating latest and
+ // latest_ref_row to the latest data prior to that immediate just past Returns true if
+ // updates were made, false if not.
+ bool advance_and_memoize(int64_t ts) {
+ // Advance the right side row index until we reach the latest right row (for each key)
+ // for the given left timestamp.
+
+ // Check if already updated for TS (or if there is no latest)
+ if (empty()) return false; // can't advance if empty
+ auto latest_time = get_latest_time();
+ if (latest_time > ts) return false; // already advanced
+
+ // Not updated. Try to update and possibly advance.
+ bool updated = false;
+ do {
+ latest_time = get_latest_time();
+ // 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 <= ts) {
+ memo_.store(get_latest_batch(), latest_ref_row_, latest_time, get_latest_key());
+ } else {
+ break; // hit a future timestamp -- done updating for now
+ }
+ updated = true;
+ } while (advance());
+ return updated;
+ }
+
+ void push(const std::shared_ptr<arrow::RecordBatch>& rb) {
+ if (rb->num_rows() > 0) {
+ queue_.push(rb);
+ } else {
+ ++batches_processed_; // don't enqueue empty batches, just record as processed
+ }
+ }
+
+ util::optional<const MemoStore::Entry*> get_memo_entry_for_key(KeyType key) {
+ auto r = memo_.get_entry_for_key(key);
+ if (r.has_value()) return r;
+ if (wildcard_key_.has_value()) r = memo_.get_entry_for_key(*wildcard_key_);
+ return r;
+ }
+
+ util::optional<int64_t> get_memo_time_for_key(KeyType key) {
+ auto r = get_memo_entry_for_key(key);
+ return r.has_value() ? util::make_optional((*r)->_time) : util::nullopt;
+ }
+
+ void remove_memo_entries_with_lesser_time(int64_t ts) {
+ memo_.remove_entries_with_lesser_time(ts);
+ }
+
+ const std::shared_ptr<Schema>& get_schema() const { return schema_; }
+
+ void set_total_batches(int n) {
+ assert(n >= 0);
+ assert(total_batches_ == -1); // shouldn't be set more than once
+ total_batches_ = n;
+ }
+
+ private:
+ // Pending record batches. The latest is the front. Batches cannot be empty.
+ ConcurrentQueue<std::shared_ptr<RecordBatch>> queue_;
+
+ // Wildcard key for this input, if applicable.
+ util::optional<KeyType> wildcard_key_;
+
+ // Schema associated with the input
+ std::shared_ptr<Schema> schema_;
+
+ // Total number of batches (only int because InputFinished uses int)
+ int total_batches_ = -1;
+
+ // Number of batches processed so far (only int because InputFinished uses int)
+ int batches_processed_ = 0;
+
+ // Index of the time col
+ col_index_t time_col_index_;
+
+ // Index of the key col
+ col_index_t key_col_index_;
+
+ // Index of the latest row reference within; if >0 then queue_ cannot be empty
+ row_index_t latest_ref_row_ =
+ 0; // must be < queue_.front()->num_rows() if queue_ is non-empty
+
+ // Stores latest known values for the various keys
+ MemoStore memo_;
+
+ // Mapping of source columns to destination columns
+ std::vector<util::optional<col_index_t>> src_to_dst_;
+};
+
+template <size_t MAX_TABLES>
+struct CompositeReferenceRow {
+ struct Entry {
+ arrow::RecordBatch* batch; // can be NULL if there's no value
+ row_index_t row;
+ };
+ Entry refs[MAX_TABLES];
+};
+
+// A table of composite reference rows. Rows maintain pointers to the
+// constituent record batches, but the overall table retains shared_ptr
+// references to ensure memory remains resident while the table is live.
+//
+// The main reason for this is that, especially for wide tables, joins
+// are effectively row-oriented, rather than column-oriented. Separating
+// the join part from the columnar materialization part simplifies the
+// logic around data types and increases efficiency.
+//
+// We don't put the shared_ptr's into the rows for efficiency reasons.
+template <size_t MAX_TABLES>
+class CompositeReferenceTable {
+ public:
+ explicit CompositeReferenceTable(size_t n_tables) : n_tables_(n_tables) {
+ assert(n_tables_ >= 1);
+ assert(n_tables_ <= MAX_TABLES);
+ }
+
+ size_t n_rows() const { return rows_.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<std::unique_ptr<InputState>>& in, int64_t tolerance) {
+ assert(in.size() == n_tables_);
+
+ // Get the LHS key
+ KeyType key = in[0]->get_latest_key();
+
+ // Add row and setup LHS
+ // (the LHS state comes just from the latest row of the LHS table)
+ assert(!in[0]->empty());
+ const std::shared_ptr<arrow::RecordBatch>& lhs_latest_batch =
+ in[0]->get_latest_batch();
+ row_index_t lhs_latest_row = in[0]->get_latest_row();
+ int64_t lhs_latest_time = in[0]->get_latest_time();
+ 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 = rows_.size() + new_batch_size;
+ if (rows_.capacity() < new_capacity) rows_.reserve(new_capacity);
+ }
+ rows_.resize(rows_.size() + 1);
+ auto& row = rows_.back();
+ row.refs[0].batch = lhs_latest_batch.get();
+ row.refs[0].row = lhs_latest_row;
+ add_record_batch_ref(lhs_latest_batch);
+
+ // 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) {
+ util::optional<const MemoStore::Entry*> opt_entry =
+ in[i]->get_memo_entry_for_key(key);
+ if (opt_entry.has_value()) {
+ assert(*opt_entry);
+ if ((*opt_entry)->_time + tolerance >= lhs_latest_time) {
+ // Have a valid entry
+ const MemoStore::Entry* entry = *opt_entry;
+ row.refs[i].batch = entry->_batch.get();
+ row.refs[i].row = entry->_row;
+ add_record_batch_ref(entry->_batch);
+ continue;
+ }
+ }
+ row.refs[i].batch = NULL;
+ row.refs[i].row = 0;
+ }
+ }
+
+ // Materializes the current reference table into a target record batch
+ Result<std::shared_ptr<RecordBatch>> materialize(
+ const std::shared_ptr<arrow::Schema>& output_schema,
+ const std::vector<std::unique_ptr<InputState>>& state) {
+ // cerr << "materialize BEGIN\n";
+ assert(state.size() == n_tables_);
+ assert(state.size() >= 1);
+
+ // Don't build empty batches
+ size_t n_rows = rows_.size();
+ if (!n_rows) return NULLPTR;
+
+ // Build the arrays column-by-column from the rows
+ std::vector<std::shared_ptr<arrow::Array>> arrays(output_schema->num_fields());
+ for (size_t i_table = 0; i_table < n_tables_; ++i_table) {
+ int n_src_cols = state.at(i_table)->get_schema()->num_fields();
+ {
+ for (col_index_t i_src_col = 0; i_src_col < n_src_cols; ++i_src_col) {
+ util::optional<col_index_t> i_dst_col_opt =
+ state[i_table]->map_src_to_dst(i_src_col);
+ if (!i_dst_col_opt) continue;
+ col_index_t i_dst_col = *i_dst_col_opt;
+ const auto& src_field = state[i_table]->get_schema()->field(i_src_col);
+ const auto& dst_field = output_schema->field(i_dst_col);
+ assert(src_field->type()->Equals(dst_field->type()));
+ assert(src_field->name() == dst_field->name());
+ const auto& field_type = src_field->type();
+
+ if (field_type->Equals(arrow::int32())) {
+ ARROW_ASSIGN_OR_RAISE(
+ arrays.at(i_dst_col),
+ (materialize_primitive_column<arrow::Int32Builder, int32_t>(i_table,
+ i_src_col)));
+ } else if (field_type->Equals(arrow::int64())) {
+ ARROW_ASSIGN_OR_RAISE(
+ arrays.at(i_dst_col),
+ (materialize_primitive_column<arrow::Int64Builder, int64_t>(i_table,
+ i_src_col)));
+ } else if (field_type->Equals(arrow::float64())) {
+ ARROW_ASSIGN_OR_RAISE(
+ arrays.at(i_dst_col),
+ (materialize_primitive_column<arrow::DoubleBuilder, double>(i_table,
+ i_src_col)));
+ } else {
+ ARROW_RETURN_NOT_OK(
+ Status::Invalid("Unsupported data type: ", src_field->name()));
+ }
+ }
+ }
+ }
+
+ // Build the result
+ assert(sizeof(size_t) >= sizeof(int64_t)); // Make takes signed int64_t for num_rows
+
+ // TODO: check n_rows for cast
+ std::shared_ptr<arrow::RecordBatch> r =
+ arrow::RecordBatch::Make(output_schema, (int64_t)n_rows, arrays);
+ return r;
+ }
+
+ // Returns true if there are no rows
+ bool empty() const { return rows_.empty(); }
+
+ private:
+ // Contains shared_ptr refs for all RecordBatches referred to by the contents of rows_
+ std::unordered_map<uintptr_t, std::shared_ptr<RecordBatch>> _ptr2ref;
+
+ // Row table references
+ std::vector<CompositeReferenceRow<MAX_TABLES>> rows_;
+
+ // Total number of tables in the composite table
+ size_t n_tables_;
+
+ // Adds a RecordBatch ref to the mapping, if needed
Review Comment:
Hmm...does anything ever remove elements from `_ptr2ref`? Will that cause a leak hanging on to all batches that contain at least one output row?
##########
cpp/src/arrow/compute/exec/options.h:
##########
@@ -361,6 +361,19 @@ class ARROW_EXPORT HashJoinNodeOptions : public ExecNodeOptions {
Expression filter;
};
+class ARROW_EXPORT AsofJoinNodeOptions : public ExecNodeOptions {
Review Comment:
If I recall correctly, this relies on `time` and `key` having the same name in all of the inputs. While this works in Acero today it will not work in Substrait (where names are erased from the plan). Instead, I think we will want to someday move to `time` being a `vector<FieldRef>` and `keys` being (regrettably) a `vector<vector<FieldRef>>` (i.e. N keys across M inputs).
##########
cpp/src/arrow/compute/exec/options.h:
##########
@@ -361,6 +361,19 @@ class ARROW_EXPORT HashJoinNodeOptions : public ExecNodeOptions {
Expression filter;
};
+class ARROW_EXPORT AsofJoinNodeOptions : public ExecNodeOptions {
+ public:
+ AsofJoinNodeOptions(FieldRef time, FieldRef keys, int64_t tolerance)
+ : time(std::move(time)), keys(std::move(keys)), tolerance(tolerance) {}
+
+ // time column
+ FieldRef time;
+ // keys used for the join. All tables must have the same join key.
+ FieldRef keys;
Review Comment:
```suggestion
// key used for the join. All tables must have the same join key.
FieldRef key;
```
If you want to future proof to multiple keys you could make this `std::vector<FieldRef>` and simply reject plans that have more than one key.
##########
cpp/src/arrow/compute/exec/options.h:
##########
@@ -361,6 +361,19 @@ class ARROW_EXPORT HashJoinNodeOptions : public ExecNodeOptions {
Expression filter;
};
+class ARROW_EXPORT AsofJoinNodeOptions : public ExecNodeOptions {
+ public:
+ AsofJoinNodeOptions(FieldRef time, FieldRef keys, int64_t tolerance)
+ : time(std::move(time)), keys(std::move(keys)), tolerance(tolerance) {}
+
+ // time column
Review Comment:
There are quite a few restrictions on these options right now. `time` must be int64? `key` must be an int64? The payload columns must meet a pretty restrictive set of types.
I think this is ok as this is an early implementation. However, we should thoroughly document these limitations as well as test for these when the node is created.
##########
cpp/src/arrow/compute/exec/asof_join_node.cc:
##########
@@ -0,0 +1,800 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+#include <iostream>
+#include <unordered_map>
+
+#include <arrow/api.h>
+#include <arrow/compute/api.h>
+#include <arrow/util/optional.h>
+#include "arrow/compute/exec/asof_join.h"
+#include "arrow/compute/exec/exec_plan.h"
+#include "arrow/compute/exec/options.h"
+#include "arrow/compute/exec/schema_util.h"
+#include "arrow/compute/exec/util.h"
+#include "arrow/result.h"
+#include "arrow/status.h"
+#include "arrow/util/checked_cast.h"
+#include "arrow/util/future.h"
+#include "arrow/util/make_unique.h"
+
+#include <condition_variable>
+#include <mutex>
+#include <thread>
+
+namespace arrow {
+namespace compute {
+
+/**
+ * Simple implementation for an unbound concurrent queue
+ */
+template <class T>
+class ConcurrentQueue {
+ public:
+ T pop() {
+ std::unique_lock<std::mutex> lock(mutex_);
+ cond_.wait(lock, [&] { return !queue_.empty(); });
+ auto item = queue_.front();
+ queue_.pop();
+ return item;
+ }
+
+ void push(const T& item) {
+ std::unique_lock<std::mutex> lock(mutex_);
+ queue_.push(item);
+ cond_.notify_one();
+ }
+
+ util::optional<T> try_pop() {
+ // Try to pop the oldest value from the queue (or return nullopt if none)
+ std::unique_lock<std::mutex> lock(mutex_);
+ if (queue_.empty()) {
+ return util::nullopt;
+ } else {
+ auto item = queue_.front();
+ queue_.pop();
+ return item;
+ }
+ }
+
+ bool empty() const {
+ std::unique_lock<std::mutex> lock(mutex_);
+ return queue_.empty();
+ }
+
+ // Un-synchronized access to front
+ // For this to be "safe":
+ // 1) the caller logically guarantees that queue is not empty
+ // 2) pop/try_pop cannot be called concurrently with this
+ const T& unsync_front() const { return queue_.front(); }
+
+ private:
+ std::queue<T> queue_;
+ mutable std::mutex mutex_;
+ std::condition_variable cond_;
+};
+
+struct MemoStore {
+ // Stores last known values for all the keys
+
+ struct Entry {
+ // Timestamp associated with the entry
+ int64_t _time;
+
+ // Batch associated with the entry (perf is probably OK for this; batches change
+ // rarely)
+ std::shared_ptr<arrow::RecordBatch> _batch;
+
+ // Row associated with the entry
+ row_index_t _row;
+ };
+
+ std::unordered_map<KeyType, Entry> _entries;
+
+ void store(const std::shared_ptr<RecordBatch>& batch, row_index_t row, int64_t time,
+ KeyType key) {
+ auto& e = _entries[key];
+ // that we can do this assignment optionally, is why we
+ // can get array 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;
+ }
+
+ util::optional<const Entry*> get_entry_for_key(KeyType key) const {
+ auto e = _entries.find(key);
+ if (_entries.end() == e) return util::nullopt;
+ return util::optional<const Entry*>(&e->second);
+ }
+
+ void remove_entries_with_lesser_time(int64_t ts) {
+ size_t dbg_size0 = _entries.size();
+ for (auto e = _entries.begin(); e != _entries.end();)
+ if (e->second._time < ts)
+ e = _entries.erase(e);
+ else
+ ++e;
+ size_t dbg_size1 = _entries.size();
+ if (dbg_size1 < dbg_size0) {
+ // cerr << "Removed " << dbg_size0-dbg_size1 << " memo entries.\n";
+ }
+ }
+};
+
+class InputState {
+ // InputState correponds to an input
+ // Input record batches are queued up in InputState until processed and
+ // turned into output record batches.
+
+ public:
+ InputState(const std::shared_ptr<arrow::Schema>& schema,
+ const std::string& time_col_name, const std::string& key_col_name,
+ util::optional<KeyType> wildcard_key)
+ : queue_(),
+ wildcard_key_(wildcard_key),
+ schema_(schema),
+ time_col_index_(
+ schema->GetFieldIndex(time_col_name)), // TODO: handle missing field name
+ key_col_index_(schema->GetFieldIndex(key_col_name)) {}
+
+ col_index_t init_src_to_dst_mapping(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 && is_time_or_key_column(i)))
+ src_to_dst_[i] = dst_offset++;
+ return dst_offset;
+ }
+
+ const util::optional<col_index_t>& map_src_to_dst(col_index_t src) const {
+ return src_to_dst_[src];
+ }
+
+ bool is_time_or_key_column(col_index_t i) const {
+ assert(i < schema_->num_fields());
+ return (i == time_col_index_) || (i == key_col_index_);
+ }
+
+ // Gets the latest row index, assuming the queue isn't empty
+ row_index_t get_latest_row() const { return latest_ref_row_; }
+
+ bool empty() const {
+ if (latest_ref_row_ > 0)
+ return false; // cannot be empty if ref row is >0 -- can avoid slow queue lock
+ // below
+ return queue_.empty();
+ }
+
+ int countbatches_processed_() const { return batches_processed_; }
+ int count_total_batches() const { return total_batches_; }
+
+ // Gets latest batch (precondition: must not be empty)
+ const std::shared_ptr<arrow::RecordBatch>& get_latest_batch() const {
+ return queue_.unsync_front();
+ }
+ KeyType get_latest_key() const {
+ return queue_.unsync_front()
+ ->column_data(key_col_index_)
+ ->GetValues<KeyType>(1)[latest_ref_row_];
+ }
+ int64_t get_latest_time() const {
+ return queue_.unsync_front()
+ ->column_data(time_col_index_)
+ ->GetValues<int64_t>(1)[latest_ref_row_];
+ }
+
+ bool finished() const { return batches_processed_ == total_batches_; }
+
+ bool advance() {
+ // 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();
+ if (have_active_batch) {
+ // If we have an active batch
+ if (++latest_ref_row_ >= (row_index_t)queue_.unsync_front()->num_rows()) {
+ // hit the end of the batch, need to get the next batch if possible.
+ ++batches_processed_;
+ latest_ref_row_ = 0;
+ have_active_batch &= !queue_.try_pop();
+ if (have_active_batch)
+ assert(queue_.unsync_front()->num_rows() > 0); // empty batches disallowed
+ }
+ }
+ return have_active_batch;
+ }
+
+ // Advance the data to be immediately past the specified TS, updating latest and
+ // latest_ref_row to the latest data prior to that immediate just past Returns true if
+ // updates were made, false if not.
+ bool advance_and_memoize(int64_t ts) {
+ // Advance the right side row index until we reach the latest right row (for each key)
+ // for the given left timestamp.
+
+ // Check if already updated for TS (or if there is no latest)
+ if (empty()) return false; // can't advance if empty
+ auto latest_time = get_latest_time();
+ if (latest_time > ts) return false; // already advanced
+
+ // Not updated. Try to update and possibly advance.
+ bool updated = false;
+ do {
+ latest_time = get_latest_time();
+ // 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 <= ts) {
+ memo_.store(get_latest_batch(), latest_ref_row_, latest_time, get_latest_key());
+ } else {
+ break; // hit a future timestamp -- done updating for now
+ }
+ updated = true;
+ } while (advance());
+ return updated;
+ }
+
+ void push(const std::shared_ptr<arrow::RecordBatch>& rb) {
+ if (rb->num_rows() > 0) {
+ queue_.push(rb);
+ } else {
+ ++batches_processed_; // don't enqueue empty batches, just record as processed
+ }
+ }
+
+ util::optional<const MemoStore::Entry*> get_memo_entry_for_key(KeyType key) {
+ auto r = memo_.get_entry_for_key(key);
+ if (r.has_value()) return r;
+ if (wildcard_key_.has_value()) r = memo_.get_entry_for_key(*wildcard_key_);
+ return r;
+ }
+
+ util::optional<int64_t> get_memo_time_for_key(KeyType key) {
+ auto r = get_memo_entry_for_key(key);
+ return r.has_value() ? util::make_optional((*r)->_time) : util::nullopt;
+ }
+
+ void remove_memo_entries_with_lesser_time(int64_t ts) {
+ memo_.remove_entries_with_lesser_time(ts);
+ }
+
+ const std::shared_ptr<Schema>& get_schema() const { return schema_; }
+
+ void set_total_batches(int n) {
+ assert(n >= 0);
+ assert(total_batches_ == -1); // shouldn't be set more than once
+ total_batches_ = n;
+ }
+
+ private:
+ // Pending record batches. The latest is the front. Batches cannot be empty.
+ ConcurrentQueue<std::shared_ptr<RecordBatch>> queue_;
+
+ // Wildcard key for this input, if applicable.
+ util::optional<KeyType> wildcard_key_;
+
+ // Schema associated with the input
+ std::shared_ptr<Schema> schema_;
+
+ // Total number of batches (only int because InputFinished uses int)
+ int total_batches_ = -1;
+
+ // Number of batches processed so far (only int because InputFinished uses int)
+ int batches_processed_ = 0;
+
+ // Index of the time col
+ col_index_t time_col_index_;
+
+ // Index of the key col
+ col_index_t key_col_index_;
+
+ // Index of the latest row reference within; if >0 then queue_ cannot be empty
+ row_index_t latest_ref_row_ =
+ 0; // must be < queue_.front()->num_rows() if queue_ is non-empty
+
+ // Stores latest known values for the various keys
+ MemoStore memo_;
+
+ // Mapping of source columns to destination columns
+ std::vector<util::optional<col_index_t>> src_to_dst_;
+};
+
+template <size_t MAX_TABLES>
+struct CompositeReferenceRow {
+ struct Entry {
+ arrow::RecordBatch* batch; // can be NULL if there's no value
+ row_index_t row;
+ };
+ Entry refs[MAX_TABLES];
+};
+
+// A table of composite reference rows. Rows maintain pointers to the
+// constituent record batches, but the overall table retains shared_ptr
+// references to ensure memory remains resident while the table is live.
+//
+// The main reason for this is that, especially for wide tables, joins
+// are effectively row-oriented, rather than column-oriented. Separating
+// the join part from the columnar materialization part simplifies the
+// logic around data types and increases efficiency.
+//
+// We don't put the shared_ptr's into the rows for efficiency reasons.
+template <size_t MAX_TABLES>
+class CompositeReferenceTable {
+ public:
+ explicit CompositeReferenceTable(size_t n_tables) : n_tables_(n_tables) {
+ assert(n_tables_ >= 1);
+ assert(n_tables_ <= MAX_TABLES);
+ }
+
+ size_t n_rows() const { return rows_.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<std::unique_ptr<InputState>>& in, int64_t tolerance) {
+ assert(in.size() == n_tables_);
+
+ // Get the LHS key
+ KeyType key = in[0]->get_latest_key();
+
+ // Add row and setup LHS
+ // (the LHS state comes just from the latest row of the LHS table)
+ assert(!in[0]->empty());
+ const std::shared_ptr<arrow::RecordBatch>& lhs_latest_batch =
+ in[0]->get_latest_batch();
+ row_index_t lhs_latest_row = in[0]->get_latest_row();
+ int64_t lhs_latest_time = in[0]->get_latest_time();
+ 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 = rows_.size() + new_batch_size;
+ if (rows_.capacity() < new_capacity) rows_.reserve(new_capacity);
+ }
+ rows_.resize(rows_.size() + 1);
+ auto& row = rows_.back();
+ row.refs[0].batch = lhs_latest_batch.get();
+ row.refs[0].row = lhs_latest_row;
+ add_record_batch_ref(lhs_latest_batch);
+
+ // 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) {
+ util::optional<const MemoStore::Entry*> opt_entry =
+ in[i]->get_memo_entry_for_key(key);
+ if (opt_entry.has_value()) {
+ assert(*opt_entry);
+ if ((*opt_entry)->_time + tolerance >= lhs_latest_time) {
+ // Have a valid entry
+ const MemoStore::Entry* entry = *opt_entry;
+ row.refs[i].batch = entry->_batch.get();
+ row.refs[i].row = entry->_row;
+ add_record_batch_ref(entry->_batch);
+ continue;
+ }
+ }
+ row.refs[i].batch = NULL;
+ row.refs[i].row = 0;
+ }
+ }
+
+ // Materializes the current reference table into a target record batch
+ Result<std::shared_ptr<RecordBatch>> materialize(
+ const std::shared_ptr<arrow::Schema>& output_schema,
+ const std::vector<std::unique_ptr<InputState>>& state) {
+ // cerr << "materialize BEGIN\n";
+ assert(state.size() == n_tables_);
+ assert(state.size() >= 1);
Review Comment:
```suggestion
DCHECK_EQ(state.size(), n_tables_);
DCHECK_GE(state.size(), 1);
```
##########
cpp/src/arrow/compute/exec/asof_join_node.cc:
##########
@@ -0,0 +1,800 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+#include <iostream>
+#include <unordered_map>
+
+#include <arrow/api.h>
+#include <arrow/compute/api.h>
+#include <arrow/util/optional.h>
+#include "arrow/compute/exec/asof_join.h"
+#include "arrow/compute/exec/exec_plan.h"
+#include "arrow/compute/exec/options.h"
+#include "arrow/compute/exec/schema_util.h"
+#include "arrow/compute/exec/util.h"
+#include "arrow/result.h"
+#include "arrow/status.h"
+#include "arrow/util/checked_cast.h"
+#include "arrow/util/future.h"
+#include "arrow/util/make_unique.h"
+
+#include <condition_variable>
+#include <mutex>
+#include <thread>
+
+namespace arrow {
+namespace compute {
+
+/**
+ * Simple implementation for an unbound concurrent queue
+ */
+template <class T>
+class ConcurrentQueue {
+ public:
+ T pop() {
+ std::unique_lock<std::mutex> lock(mutex_);
+ cond_.wait(lock, [&] { return !queue_.empty(); });
+ auto item = queue_.front();
+ queue_.pop();
+ return item;
+ }
+
+ void push(const T& item) {
+ std::unique_lock<std::mutex> lock(mutex_);
+ queue_.push(item);
+ cond_.notify_one();
+ }
+
+ util::optional<T> try_pop() {
+ // Try to pop the oldest value from the queue (or return nullopt if none)
+ std::unique_lock<std::mutex> lock(mutex_);
+ if (queue_.empty()) {
+ return util::nullopt;
+ } else {
+ auto item = queue_.front();
+ queue_.pop();
+ return item;
+ }
+ }
+
+ bool empty() const {
+ std::unique_lock<std::mutex> lock(mutex_);
+ return queue_.empty();
+ }
+
+ // Un-synchronized access to front
+ // For this to be "safe":
+ // 1) the caller logically guarantees that queue is not empty
+ // 2) pop/try_pop cannot be called concurrently with this
+ const T& unsync_front() const { return queue_.front(); }
+
+ private:
+ std::queue<T> queue_;
+ mutable std::mutex mutex_;
+ std::condition_variable cond_;
+};
+
+struct MemoStore {
+ // Stores last known values for all the keys
+
+ struct Entry {
+ // Timestamp associated with the entry
+ int64_t _time;
+
+ // Batch associated with the entry (perf is probably OK for this; batches change
+ // rarely)
+ std::shared_ptr<arrow::RecordBatch> _batch;
+
+ // Row associated with the entry
+ row_index_t _row;
+ };
+
+ std::unordered_map<KeyType, Entry> _entries;
+
+ void store(const std::shared_ptr<RecordBatch>& batch, row_index_t row, int64_t time,
+ KeyType key) {
+ auto& e = _entries[key];
+ // that we can do this assignment optionally, is why we
+ // can get array 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;
+ }
+
+ util::optional<const Entry*> get_entry_for_key(KeyType key) const {
+ auto e = _entries.find(key);
+ if (_entries.end() == e) return util::nullopt;
+ return util::optional<const Entry*>(&e->second);
+ }
+
+ void remove_entries_with_lesser_time(int64_t ts) {
+ size_t dbg_size0 = _entries.size();
+ for (auto e = _entries.begin(); e != _entries.end();)
+ if (e->second._time < ts)
+ e = _entries.erase(e);
+ else
+ ++e;
+ size_t dbg_size1 = _entries.size();
+ if (dbg_size1 < dbg_size0) {
+ // cerr << "Removed " << dbg_size0-dbg_size1 << " memo entries.\n";
+ }
+ }
+};
+
+class InputState {
+ // InputState correponds to an input
+ // Input record batches are queued up in InputState until processed and
+ // turned into output record batches.
+
+ public:
+ InputState(const std::shared_ptr<arrow::Schema>& schema,
+ const std::string& time_col_name, const std::string& key_col_name,
+ util::optional<KeyType> wildcard_key)
+ : queue_(),
+ wildcard_key_(wildcard_key),
+ schema_(schema),
+ time_col_index_(
+ schema->GetFieldIndex(time_col_name)), // TODO: handle missing field name
+ key_col_index_(schema->GetFieldIndex(key_col_name)) {}
+
+ col_index_t init_src_to_dst_mapping(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 && is_time_or_key_column(i)))
+ src_to_dst_[i] = dst_offset++;
+ return dst_offset;
+ }
+
+ const util::optional<col_index_t>& map_src_to_dst(col_index_t src) const {
+ return src_to_dst_[src];
+ }
+
+ bool is_time_or_key_column(col_index_t i) const {
+ assert(i < schema_->num_fields());
+ return (i == time_col_index_) || (i == key_col_index_);
+ }
+
+ // Gets the latest row index, assuming the queue isn't empty
+ row_index_t get_latest_row() const { return latest_ref_row_; }
+
+ bool empty() const {
+ if (latest_ref_row_ > 0)
+ return false; // cannot be empty if ref row is >0 -- can avoid slow queue lock
+ // below
+ return queue_.empty();
+ }
+
+ int countbatches_processed_() const { return batches_processed_; }
+ int count_total_batches() const { return total_batches_; }
+
+ // Gets latest batch (precondition: must not be empty)
+ const std::shared_ptr<arrow::RecordBatch>& get_latest_batch() const {
+ return queue_.unsync_front();
+ }
+ KeyType get_latest_key() const {
+ return queue_.unsync_front()
+ ->column_data(key_col_index_)
+ ->GetValues<KeyType>(1)[latest_ref_row_];
+ }
+ int64_t get_latest_time() const {
+ return queue_.unsync_front()
+ ->column_data(time_col_index_)
+ ->GetValues<int64_t>(1)[latest_ref_row_];
+ }
+
+ bool finished() const { return batches_processed_ == total_batches_; }
+
+ bool advance() {
+ // 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();
+ if (have_active_batch) {
+ // If we have an active batch
+ if (++latest_ref_row_ >= (row_index_t)queue_.unsync_front()->num_rows()) {
+ // hit the end of the batch, need to get the next batch if possible.
+ ++batches_processed_;
+ latest_ref_row_ = 0;
+ have_active_batch &= !queue_.try_pop();
+ if (have_active_batch)
+ assert(queue_.unsync_front()->num_rows() > 0); // empty batches disallowed
+ }
+ }
+ return have_active_batch;
+ }
+
+ // Advance the data to be immediately past the specified TS, updating latest and
+ // latest_ref_row to the latest data prior to that immediate just past Returns true if
+ // updates were made, false if not.
+ bool advance_and_memoize(int64_t ts) {
+ // Advance the right side row index until we reach the latest right row (for each key)
+ // for the given left timestamp.
+
+ // Check if already updated for TS (or if there is no latest)
+ if (empty()) return false; // can't advance if empty
+ auto latest_time = get_latest_time();
+ if (latest_time > ts) return false; // already advanced
+
+ // Not updated. Try to update and possibly advance.
+ bool updated = false;
+ do {
+ latest_time = get_latest_time();
+ // 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 <= ts) {
+ memo_.store(get_latest_batch(), latest_ref_row_, latest_time, get_latest_key());
+ } else {
+ break; // hit a future timestamp -- done updating for now
+ }
+ updated = true;
+ } while (advance());
+ return updated;
+ }
+
+ void push(const std::shared_ptr<arrow::RecordBatch>& rb) {
+ if (rb->num_rows() > 0) {
+ queue_.push(rb);
+ } else {
+ ++batches_processed_; // don't enqueue empty batches, just record as processed
+ }
+ }
+
+ util::optional<const MemoStore::Entry*> get_memo_entry_for_key(KeyType key) {
+ auto r = memo_.get_entry_for_key(key);
+ if (r.has_value()) return r;
+ if (wildcard_key_.has_value()) r = memo_.get_entry_for_key(*wildcard_key_);
+ return r;
+ }
+
+ util::optional<int64_t> get_memo_time_for_key(KeyType key) {
+ auto r = get_memo_entry_for_key(key);
+ return r.has_value() ? util::make_optional((*r)->_time) : util::nullopt;
+ }
+
+ void remove_memo_entries_with_lesser_time(int64_t ts) {
+ memo_.remove_entries_with_lesser_time(ts);
+ }
+
+ const std::shared_ptr<Schema>& get_schema() const { return schema_; }
+
+ void set_total_batches(int n) {
+ assert(n >= 0);
+ assert(total_batches_ == -1); // shouldn't be set more than once
+ total_batches_ = n;
+ }
+
+ private:
+ // Pending record batches. The latest is the front. Batches cannot be empty.
+ ConcurrentQueue<std::shared_ptr<RecordBatch>> queue_;
+
+ // Wildcard key for this input, if applicable.
+ util::optional<KeyType> wildcard_key_;
+
+ // Schema associated with the input
+ std::shared_ptr<Schema> schema_;
+
+ // Total number of batches (only int because InputFinished uses int)
+ int total_batches_ = -1;
+
+ // Number of batches processed so far (only int because InputFinished uses int)
+ int batches_processed_ = 0;
+
+ // Index of the time col
+ col_index_t time_col_index_;
+
+ // Index of the key col
+ col_index_t key_col_index_;
+
+ // Index of the latest row reference within; if >0 then queue_ cannot be empty
+ row_index_t latest_ref_row_ =
+ 0; // must be < queue_.front()->num_rows() if queue_ is non-empty
+
+ // Stores latest known values for the various keys
+ MemoStore memo_;
+
+ // Mapping of source columns to destination columns
+ std::vector<util::optional<col_index_t>> src_to_dst_;
+};
+
+template <size_t MAX_TABLES>
+struct CompositeReferenceRow {
+ struct Entry {
+ arrow::RecordBatch* batch; // can be NULL if there's no value
+ row_index_t row;
+ };
+ Entry refs[MAX_TABLES];
+};
+
+// A table of composite reference rows. Rows maintain pointers to the
+// constituent record batches, but the overall table retains shared_ptr
+// references to ensure memory remains resident while the table is live.
+//
+// The main reason for this is that, especially for wide tables, joins
+// are effectively row-oriented, rather than column-oriented. Separating
+// the join part from the columnar materialization part simplifies the
+// logic around data types and increases efficiency.
+//
+// We don't put the shared_ptr's into the rows for efficiency reasons.
+template <size_t MAX_TABLES>
+class CompositeReferenceTable {
+ public:
+ explicit CompositeReferenceTable(size_t n_tables) : n_tables_(n_tables) {
+ assert(n_tables_ >= 1);
+ assert(n_tables_ <= MAX_TABLES);
+ }
+
+ size_t n_rows() const { return rows_.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<std::unique_ptr<InputState>>& in, int64_t tolerance) {
+ assert(in.size() == n_tables_);
+
+ // Get the LHS key
+ KeyType key = in[0]->get_latest_key();
+
+ // Add row and setup LHS
+ // (the LHS state comes just from the latest row of the LHS table)
+ assert(!in[0]->empty());
+ const std::shared_ptr<arrow::RecordBatch>& lhs_latest_batch =
+ in[0]->get_latest_batch();
+ row_index_t lhs_latest_row = in[0]->get_latest_row();
+ int64_t lhs_latest_time = in[0]->get_latest_time();
+ 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 = rows_.size() + new_batch_size;
+ if (rows_.capacity() < new_capacity) rows_.reserve(new_capacity);
+ }
+ rows_.resize(rows_.size() + 1);
+ auto& row = rows_.back();
+ row.refs[0].batch = lhs_latest_batch.get();
+ row.refs[0].row = lhs_latest_row;
+ add_record_batch_ref(lhs_latest_batch);
+
+ // 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) {
+ util::optional<const MemoStore::Entry*> opt_entry =
+ in[i]->get_memo_entry_for_key(key);
+ if (opt_entry.has_value()) {
+ assert(*opt_entry);
+ if ((*opt_entry)->_time + tolerance >= lhs_latest_time) {
+ // Have a valid entry
+ const MemoStore::Entry* entry = *opt_entry;
+ row.refs[i].batch = entry->_batch.get();
+ row.refs[i].row = entry->_row;
+ add_record_batch_ref(entry->_batch);
+ continue;
+ }
+ }
+ row.refs[i].batch = NULL;
+ row.refs[i].row = 0;
+ }
+ }
+
+ // Materializes the current reference table into a target record batch
+ Result<std::shared_ptr<RecordBatch>> materialize(
+ const std::shared_ptr<arrow::Schema>& output_schema,
+ const std::vector<std::unique_ptr<InputState>>& state) {
+ // cerr << "materialize BEGIN\n";
+ assert(state.size() == n_tables_);
+ assert(state.size() >= 1);
+
+ // Don't build empty batches
+ size_t n_rows = rows_.size();
+ if (!n_rows) return NULLPTR;
+
+ // Build the arrays column-by-column from the rows
+ std::vector<std::shared_ptr<arrow::Array>> arrays(output_schema->num_fields());
+ for (size_t i_table = 0; i_table < n_tables_; ++i_table) {
+ int n_src_cols = state.at(i_table)->get_schema()->num_fields();
+ {
+ for (col_index_t i_src_col = 0; i_src_col < n_src_cols; ++i_src_col) {
+ util::optional<col_index_t> i_dst_col_opt =
+ state[i_table]->map_src_to_dst(i_src_col);
+ if (!i_dst_col_opt) continue;
+ col_index_t i_dst_col = *i_dst_col_opt;
+ const auto& src_field = state[i_table]->get_schema()->field(i_src_col);
+ const auto& dst_field = output_schema->field(i_dst_col);
+ assert(src_field->type()->Equals(dst_field->type()));
+ assert(src_field->name() == dst_field->name());
+ const auto& field_type = src_field->type();
+
+ if (field_type->Equals(arrow::int32())) {
+ ARROW_ASSIGN_OR_RAISE(
+ arrays.at(i_dst_col),
+ (materialize_primitive_column<arrow::Int32Builder, int32_t>(i_table,
+ i_src_col)));
+ } else if (field_type->Equals(arrow::int64())) {
+ ARROW_ASSIGN_OR_RAISE(
+ arrays.at(i_dst_col),
+ (materialize_primitive_column<arrow::Int64Builder, int64_t>(i_table,
+ i_src_col)));
+ } else if (field_type->Equals(arrow::float64())) {
+ ARROW_ASSIGN_OR_RAISE(
+ arrays.at(i_dst_col),
+ (materialize_primitive_column<arrow::DoubleBuilder, double>(i_table,
+ i_src_col)));
+ } else {
+ ARROW_RETURN_NOT_OK(
+ Status::Invalid("Unsupported data type: ", src_field->name()));
+ }
+ }
+ }
+ }
+
+ // Build the result
+ assert(sizeof(size_t) >= sizeof(int64_t)); // Make takes signed int64_t for num_rows
+
+ // TODO: check n_rows for cast
+ std::shared_ptr<arrow::RecordBatch> r =
+ arrow::RecordBatch::Make(output_schema, (int64_t)n_rows, arrays);
+ return r;
+ }
+
+ // Returns true if there are no rows
+ bool empty() const { return rows_.empty(); }
+
+ private:
+ // Contains shared_ptr refs for all RecordBatches referred to by the contents of rows_
+ std::unordered_map<uintptr_t, std::shared_ptr<RecordBatch>> _ptr2ref;
+
+ // Row table references
+ std::vector<CompositeReferenceRow<MAX_TABLES>> rows_;
+
+ // Total number of tables in the composite table
+ size_t n_tables_;
+
+ // Adds a RecordBatch ref to the mapping, if needed
+ void add_record_batch_ref(const std::shared_ptr<RecordBatch>& ref) {
+ if (!_ptr2ref.count((uintptr_t)ref.get())) _ptr2ref[(uintptr_t)ref.get()] = ref;
+ }
+
+ template <class Builder, class PrimitiveType>
+ Result<std::shared_ptr<Array>> materialize_primitive_column(size_t i_table,
+ col_index_t i_col) {
+ Builder builder;
+ ARROW_RETURN_NOT_OK(builder.Reserve(rows_.size()));
+ for (row_index_t i_row = 0; i_row < rows_.size(); ++i_row) {
+ const auto& ref = rows_[i_row].refs[i_table];
+ if (ref.batch) {
+ builder.UnsafeAppend(
+ ref.batch->column_data(i_col)->template GetValues<PrimitiveType>(1)[ref.row]);
+ } else {
+ builder.UnsafeAppendNull();
+ }
+ }
+ std::shared_ptr<Array> result;
+ ARROW_RETURN_NOT_OK(builder.Finish(&result));
+ return result;
+ }
+};
+
+class AsofJoinNode : public ExecNode {
+ // Constructs labels for inputs
+ static std::vector<std::string> build_input_labels(
+ const std::vector<ExecNode*>& inputs) {
+ std::vector<std::string> r(inputs.size());
+ for (size_t i = 0; i < r.size(); ++i) r[i] = "input_" + std::to_string(i) + "_label";
+ return r;
+ }
Review Comment:
Is this used?
##########
cpp/src/arrow/compute/exec/asof_join_node.cc:
##########
@@ -0,0 +1,800 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+#include <iostream>
+#include <unordered_map>
+
+#include <arrow/api.h>
+#include <arrow/compute/api.h>
+#include <arrow/util/optional.h>
+#include "arrow/compute/exec/asof_join.h"
+#include "arrow/compute/exec/exec_plan.h"
+#include "arrow/compute/exec/options.h"
+#include "arrow/compute/exec/schema_util.h"
+#include "arrow/compute/exec/util.h"
+#include "arrow/result.h"
+#include "arrow/status.h"
+#include "arrow/util/checked_cast.h"
+#include "arrow/util/future.h"
+#include "arrow/util/make_unique.h"
+
+#include <condition_variable>
+#include <mutex>
+#include <thread>
+
+namespace arrow {
+namespace compute {
+
+/**
+ * Simple implementation for an unbound concurrent queue
+ */
+template <class T>
+class ConcurrentQueue {
+ public:
+ T pop() {
+ std::unique_lock<std::mutex> lock(mutex_);
+ cond_.wait(lock, [&] { return !queue_.empty(); });
+ auto item = queue_.front();
+ queue_.pop();
+ return item;
+ }
+
+ void push(const T& item) {
+ std::unique_lock<std::mutex> lock(mutex_);
+ queue_.push(item);
+ cond_.notify_one();
+ }
+
+ util::optional<T> try_pop() {
+ // Try to pop the oldest value from the queue (or return nullopt if none)
+ std::unique_lock<std::mutex> lock(mutex_);
+ if (queue_.empty()) {
+ return util::nullopt;
+ } else {
+ auto item = queue_.front();
+ queue_.pop();
+ return item;
+ }
+ }
+
+ bool empty() const {
+ std::unique_lock<std::mutex> lock(mutex_);
+ return queue_.empty();
+ }
+
+ // Un-synchronized access to front
+ // For this to be "safe":
+ // 1) the caller logically guarantees that queue is not empty
+ // 2) pop/try_pop cannot be called concurrently with this
+ const T& unsync_front() const { return queue_.front(); }
+
+ private:
+ std::queue<T> queue_;
+ mutable std::mutex mutex_;
+ std::condition_variable cond_;
+};
+
+struct MemoStore {
+ // Stores last known values for all the keys
+
+ struct Entry {
+ // Timestamp associated with the entry
+ int64_t _time;
+
+ // Batch associated with the entry (perf is probably OK for this; batches change
+ // rarely)
+ std::shared_ptr<arrow::RecordBatch> _batch;
+
+ // Row associated with the entry
+ row_index_t _row;
+ };
+
+ std::unordered_map<KeyType, Entry> _entries;
+
+ void store(const std::shared_ptr<RecordBatch>& batch, row_index_t row, int64_t time,
+ KeyType key) {
+ auto& e = _entries[key];
+ // that we can do this assignment optionally, is why we
+ // can get array 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;
+ }
+
+ util::optional<const Entry*> get_entry_for_key(KeyType key) const {
+ auto e = _entries.find(key);
+ if (_entries.end() == e) return util::nullopt;
+ return util::optional<const Entry*>(&e->second);
+ }
+
+ void remove_entries_with_lesser_time(int64_t ts) {
+ size_t dbg_size0 = _entries.size();
+ for (auto e = _entries.begin(); e != _entries.end();)
+ if (e->second._time < ts)
+ e = _entries.erase(e);
+ else
+ ++e;
+ size_t dbg_size1 = _entries.size();
+ if (dbg_size1 < dbg_size0) {
+ // cerr << "Removed " << dbg_size0-dbg_size1 << " memo entries.\n";
+ }
+ }
+};
+
+class InputState {
+ // InputState correponds to an input
+ // Input record batches are queued up in InputState until processed and
+ // turned into output record batches.
+
+ public:
+ InputState(const std::shared_ptr<arrow::Schema>& schema,
+ const std::string& time_col_name, const std::string& key_col_name,
+ util::optional<KeyType> wildcard_key)
+ : queue_(),
+ wildcard_key_(wildcard_key),
+ schema_(schema),
+ time_col_index_(
+ schema->GetFieldIndex(time_col_name)), // TODO: handle missing field name
+ key_col_index_(schema->GetFieldIndex(key_col_name)) {}
+
+ col_index_t init_src_to_dst_mapping(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 && is_time_or_key_column(i)))
+ src_to_dst_[i] = dst_offset++;
+ return dst_offset;
+ }
+
+ const util::optional<col_index_t>& map_src_to_dst(col_index_t src) const {
+ return src_to_dst_[src];
+ }
+
+ bool is_time_or_key_column(col_index_t i) const {
+ assert(i < schema_->num_fields());
+ return (i == time_col_index_) || (i == key_col_index_);
+ }
+
+ // Gets the latest row index, assuming the queue isn't empty
+ row_index_t get_latest_row() const { return latest_ref_row_; }
+
+ bool empty() const {
+ if (latest_ref_row_ > 0)
+ return false; // cannot be empty if ref row is >0 -- can avoid slow queue lock
+ // below
+ return queue_.empty();
+ }
+
+ int countbatches_processed_() const { return batches_processed_; }
+ int count_total_batches() const { return total_batches_; }
+
+ // Gets latest batch (precondition: must not be empty)
+ const std::shared_ptr<arrow::RecordBatch>& get_latest_batch() const {
+ return queue_.unsync_front();
+ }
+ KeyType get_latest_key() const {
+ return queue_.unsync_front()
+ ->column_data(key_col_index_)
+ ->GetValues<KeyType>(1)[latest_ref_row_];
+ }
+ int64_t get_latest_time() const {
+ return queue_.unsync_front()
+ ->column_data(time_col_index_)
+ ->GetValues<int64_t>(1)[latest_ref_row_];
+ }
+
+ bool finished() const { return batches_processed_ == total_batches_; }
+
+ bool advance() {
+ // 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();
+ if (have_active_batch) {
+ // If we have an active batch
+ if (++latest_ref_row_ >= (row_index_t)queue_.unsync_front()->num_rows()) {
+ // hit the end of the batch, need to get the next batch if possible.
+ ++batches_processed_;
+ latest_ref_row_ = 0;
+ have_active_batch &= !queue_.try_pop();
+ if (have_active_batch)
+ assert(queue_.unsync_front()->num_rows() > 0); // empty batches disallowed
+ }
+ }
+ return have_active_batch;
+ }
+
+ // Advance the data to be immediately past the specified TS, updating latest and
+ // latest_ref_row to the latest data prior to that immediate just past Returns true if
+ // updates were made, false if not.
+ bool advance_and_memoize(int64_t ts) {
+ // Advance the right side row index until we reach the latest right row (for each key)
+ // for the given left timestamp.
+
+ // Check if already updated for TS (or if there is no latest)
+ if (empty()) return false; // can't advance if empty
+ auto latest_time = get_latest_time();
+ if (latest_time > ts) return false; // already advanced
+
+ // Not updated. Try to update and possibly advance.
+ bool updated = false;
+ do {
+ latest_time = get_latest_time();
+ // 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 <= ts) {
+ memo_.store(get_latest_batch(), latest_ref_row_, latest_time, get_latest_key());
+ } else {
+ break; // hit a future timestamp -- done updating for now
+ }
+ updated = true;
+ } while (advance());
+ return updated;
+ }
+
+ void push(const std::shared_ptr<arrow::RecordBatch>& rb) {
+ if (rb->num_rows() > 0) {
+ queue_.push(rb);
+ } else {
+ ++batches_processed_; // don't enqueue empty batches, just record as processed
+ }
+ }
+
+ util::optional<const MemoStore::Entry*> get_memo_entry_for_key(KeyType key) {
+ auto r = memo_.get_entry_for_key(key);
+ if (r.has_value()) return r;
+ if (wildcard_key_.has_value()) r = memo_.get_entry_for_key(*wildcard_key_);
+ return r;
+ }
+
+ util::optional<int64_t> get_memo_time_for_key(KeyType key) {
+ auto r = get_memo_entry_for_key(key);
+ return r.has_value() ? util::make_optional((*r)->_time) : util::nullopt;
+ }
+
+ void remove_memo_entries_with_lesser_time(int64_t ts) {
+ memo_.remove_entries_with_lesser_time(ts);
+ }
+
+ const std::shared_ptr<Schema>& get_schema() const { return schema_; }
+
+ void set_total_batches(int n) {
+ assert(n >= 0);
+ assert(total_batches_ == -1); // shouldn't be set more than once
+ total_batches_ = n;
+ }
+
+ private:
+ // Pending record batches. The latest is the front. Batches cannot be empty.
+ ConcurrentQueue<std::shared_ptr<RecordBatch>> queue_;
+
+ // Wildcard key for this input, if applicable.
+ util::optional<KeyType> wildcard_key_;
+
+ // Schema associated with the input
+ std::shared_ptr<Schema> schema_;
+
+ // Total number of batches (only int because InputFinished uses int)
+ int total_batches_ = -1;
+
+ // Number of batches processed so far (only int because InputFinished uses int)
+ int batches_processed_ = 0;
+
+ // Index of the time col
+ col_index_t time_col_index_;
+
+ // Index of the key col
+ col_index_t key_col_index_;
+
+ // Index of the latest row reference within; if >0 then queue_ cannot be empty
+ row_index_t latest_ref_row_ =
+ 0; // must be < queue_.front()->num_rows() if queue_ is non-empty
+
+ // Stores latest known values for the various keys
+ MemoStore memo_;
+
+ // Mapping of source columns to destination columns
+ std::vector<util::optional<col_index_t>> src_to_dst_;
+};
+
+template <size_t MAX_TABLES>
+struct CompositeReferenceRow {
+ struct Entry {
+ arrow::RecordBatch* batch; // can be NULL if there's no value
+ row_index_t row;
+ };
+ Entry refs[MAX_TABLES];
+};
+
+// A table of composite reference rows. Rows maintain pointers to the
+// constituent record batches, but the overall table retains shared_ptr
+// references to ensure memory remains resident while the table is live.
+//
+// The main reason for this is that, especially for wide tables, joins
+// are effectively row-oriented, rather than column-oriented. Separating
+// the join part from the columnar materialization part simplifies the
+// logic around data types and increases efficiency.
+//
+// We don't put the shared_ptr's into the rows for efficiency reasons.
+template <size_t MAX_TABLES>
+class CompositeReferenceTable {
+ public:
+ explicit CompositeReferenceTable(size_t n_tables) : n_tables_(n_tables) {
+ assert(n_tables_ >= 1);
+ assert(n_tables_ <= MAX_TABLES);
+ }
+
+ size_t n_rows() const { return rows_.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<std::unique_ptr<InputState>>& in, int64_t tolerance) {
+ assert(in.size() == n_tables_);
+
+ // Get the LHS key
+ KeyType key = in[0]->get_latest_key();
+
+ // Add row and setup LHS
+ // (the LHS state comes just from the latest row of the LHS table)
+ assert(!in[0]->empty());
+ const std::shared_ptr<arrow::RecordBatch>& lhs_latest_batch =
+ in[0]->get_latest_batch();
+ row_index_t lhs_latest_row = in[0]->get_latest_row();
+ int64_t lhs_latest_time = in[0]->get_latest_time();
+ 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 = rows_.size() + new_batch_size;
+ if (rows_.capacity() < new_capacity) rows_.reserve(new_capacity);
+ }
+ rows_.resize(rows_.size() + 1);
+ auto& row = rows_.back();
+ row.refs[0].batch = lhs_latest_batch.get();
+ row.refs[0].row = lhs_latest_row;
+ add_record_batch_ref(lhs_latest_batch);
+
+ // 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) {
+ util::optional<const MemoStore::Entry*> opt_entry =
+ in[i]->get_memo_entry_for_key(key);
+ if (opt_entry.has_value()) {
+ assert(*opt_entry);
+ if ((*opt_entry)->_time + tolerance >= lhs_latest_time) {
+ // Have a valid entry
+ const MemoStore::Entry* entry = *opt_entry;
+ row.refs[i].batch = entry->_batch.get();
+ row.refs[i].row = entry->_row;
+ add_record_batch_ref(entry->_batch);
+ continue;
+ }
+ }
+ row.refs[i].batch = NULL;
+ row.refs[i].row = 0;
+ }
+ }
+
+ // Materializes the current reference table into a target record batch
+ Result<std::shared_ptr<RecordBatch>> materialize(
+ const std::shared_ptr<arrow::Schema>& output_schema,
+ const std::vector<std::unique_ptr<InputState>>& state) {
+ // cerr << "materialize BEGIN\n";
+ assert(state.size() == n_tables_);
+ assert(state.size() >= 1);
+
+ // Don't build empty batches
+ size_t n_rows = rows_.size();
+ if (!n_rows) return NULLPTR;
+
+ // Build the arrays column-by-column from the rows
+ std::vector<std::shared_ptr<arrow::Array>> arrays(output_schema->num_fields());
+ for (size_t i_table = 0; i_table < n_tables_; ++i_table) {
+ int n_src_cols = state.at(i_table)->get_schema()->num_fields();
+ {
+ for (col_index_t i_src_col = 0; i_src_col < n_src_cols; ++i_src_col) {
+ util::optional<col_index_t> i_dst_col_opt =
+ state[i_table]->map_src_to_dst(i_src_col);
+ if (!i_dst_col_opt) continue;
+ col_index_t i_dst_col = *i_dst_col_opt;
+ const auto& src_field = state[i_table]->get_schema()->field(i_src_col);
+ const auto& dst_field = output_schema->field(i_dst_col);
+ assert(src_field->type()->Equals(dst_field->type()));
+ assert(src_field->name() == dst_field->name());
+ const auto& field_type = src_field->type();
+
+ if (field_type->Equals(arrow::int32())) {
+ ARROW_ASSIGN_OR_RAISE(
+ arrays.at(i_dst_col),
+ (materialize_primitive_column<arrow::Int32Builder, int32_t>(i_table,
+ i_src_col)));
+ } else if (field_type->Equals(arrow::int64())) {
+ ARROW_ASSIGN_OR_RAISE(
+ arrays.at(i_dst_col),
+ (materialize_primitive_column<arrow::Int64Builder, int64_t>(i_table,
+ i_src_col)));
+ } else if (field_type->Equals(arrow::float64())) {
+ ARROW_ASSIGN_OR_RAISE(
+ arrays.at(i_dst_col),
+ (materialize_primitive_column<arrow::DoubleBuilder, double>(i_table,
+ i_src_col)));
+ } else {
+ ARROW_RETURN_NOT_OK(
+ Status::Invalid("Unsupported data type: ", src_field->name()));
+ }
+ }
+ }
+ }
+
+ // Build the result
+ assert(sizeof(size_t) >= sizeof(int64_t)); // Make takes signed int64_t for num_rows
+
+ // TODO: check n_rows for cast
+ std::shared_ptr<arrow::RecordBatch> r =
+ arrow::RecordBatch::Make(output_schema, (int64_t)n_rows, arrays);
+ return r;
+ }
+
+ // Returns true if there are no rows
+ bool empty() const { return rows_.empty(); }
+
+ private:
+ // Contains shared_ptr refs for all RecordBatches referred to by the contents of rows_
+ std::unordered_map<uintptr_t, std::shared_ptr<RecordBatch>> _ptr2ref;
+
+ // Row table references
+ std::vector<CompositeReferenceRow<MAX_TABLES>> rows_;
+
+ // Total number of tables in the composite table
+ size_t n_tables_;
+
+ // Adds a RecordBatch ref to the mapping, if needed
+ void add_record_batch_ref(const std::shared_ptr<RecordBatch>& ref) {
+ if (!_ptr2ref.count((uintptr_t)ref.get())) _ptr2ref[(uintptr_t)ref.get()] = ref;
+ }
+
+ template <class Builder, class PrimitiveType>
+ Result<std::shared_ptr<Array>> materialize_primitive_column(size_t i_table,
+ col_index_t i_col) {
+ Builder builder;
+ ARROW_RETURN_NOT_OK(builder.Reserve(rows_.size()));
+ for (row_index_t i_row = 0; i_row < rows_.size(); ++i_row) {
+ const auto& ref = rows_[i_row].refs[i_table];
+ if (ref.batch) {
+ builder.UnsafeAppend(
+ ref.batch->column_data(i_col)->template GetValues<PrimitiveType>(1)[ref.row]);
+ } else {
+ builder.UnsafeAppendNull();
+ }
+ }
+ std::shared_ptr<Array> result;
+ ARROW_RETURN_NOT_OK(builder.Finish(&result));
+ return result;
+ }
+};
+
+class AsofJoinNode : public ExecNode {
+ // Constructs labels for inputs
+ static std::vector<std::string> build_input_labels(
+ const std::vector<ExecNode*>& inputs) {
+ std::vector<std::string> r(inputs.size());
+ for (size_t i = 0; i < r.size(); ++i) r[i] = "input_" + std::to_string(i) + "_label";
+ return r;
+ }
+
+ // Advances the RHS as far as possible to be up to date for the current LHS timestamp
+ bool update_rhs() {
+ auto& lhs = *_state.at(0);
+ auto lhs_latest_time = lhs.get_latest_time();
+ bool any_updated = false;
+ for (size_t i = 1; i < _state.size(); ++i)
+ any_updated |= _state[i]->advance_and_memoize(lhs_latest_time);
+ return any_updated;
+ }
+
+ // Returns false if RHS not up to date for LHS
+ bool is_up_to_date_for_lhs_row() const {
+ auto& lhs = *_state[0];
+ if (lhs.empty()) return false; // can't proceed if nothing on the LHS
+ int64_t lhs_ts = lhs.get_latest_time();
+ for (size_t i = 1; i < _state.size(); ++i) {
+ auto& rhs = *_state[i];
+ if (!rhs.finished()) {
+ // If RHS is finished, then we know it's up to date (but if it isn't, it might be
+ // up to date)
+ if (rhs.empty())
+ return false; // RHS isn't finished, but is empty --> not up to date
+ if (lhs_ts >= rhs.get_latest_time())
+ return false; // TS not up to date (and not finished)
+ }
+ }
+ return true;
+ }
+
+ Result<std::shared_ptr<RecordBatch>> process_inner() {
+ assert(!_state.empty());
+ auto& lhs = *_state.at(0);
+
+ // Construct new target table if needed
+ CompositeReferenceTable<MAX_JOIN_TABLES> dst(_state.size());
+
+ // 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;
+
+ // Advance each of the RHS as far as possible to be up to date for the LHS timestamp
+ bool any_advanced = update_rhs();
+
+ // Only update if we have up-to-date information for the LHS row
+ if (is_up_to_date_for_lhs_row()) {
+ dst.emplace(_state, _options.tolerance);
+ if (!lhs.advance()) break; // if we can't advance LHS, we're done for this batch
+ } else {
+ if ((!any_advanced) && (_state.size() > 1)) break; // need to wait for new data
+ }
+ }
+
+ // Prune memo entries that have expired (to bound memory consumption)
+ if (!lhs.empty()) {
+ for (size_t i = 1; i < _state.size(); ++i) {
+ _state[i]->remove_memo_entries_with_lesser_time(lhs.get_latest_time() -
+ _options.tolerance);
+ }
+ }
+
+ // Emit the batch
+ if (dst.empty()) {
+ return NULLPTR;
+ } else {
+ return dst.materialize(output_schema(), _state);
+ }
+ }
+
+ void process() {
+ std::cerr << "process() begin\n";
+
+ std::lock_guard<std::mutex> guard(_gate);
+ if (finished_.is_finished()) {
+ std::cerr << "InputReceived EARLYEND\n";
+ return;
+ }
+
+ // Process batches while we have data
+ for (;;) {
+ Result<std::shared_ptr<RecordBatch>> result = process_inner();
+
+ if (result.ok()) {
+ auto out_rb = *result;
+ if (!out_rb) break;
+ ++_progress_batches_produced;
+ ExecBatch out_b(*out_rb);
+ outputs_[0]->InputReceived(this, std::move(out_b));
+ } else {
+ StopProducing();
+ ErrorIfNotOk(result.status());
+ return;
+ }
+ }
+
+ std::cerr << "process() end\n";
+
+ // 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 (_state.at(0)->finished()) {
+ total_batches_produced_ = util::make_optional<int>(_progress_batches_produced);
+ StopProducing();
+ assert(total_batches_produced_.has_value());
+ outputs_[0]->InputFinished(this, *total_batches_produced_);
+ }
+ }
+
+ void process_thread() {
+ std::cerr << "AsofJoinNode::process_thread started.\n";
+ for (;;) {
+ if (!_process.pop()) {
+ std::cerr << "AsofJoinNode::process_thread done.\n";
+ return;
+ }
+ process();
+ }
+ }
+
+ static void process_thread_wrapper(AsofJoinNode* node) { node->process_thread(); }
+
+ public:
+ AsofJoinNode(ExecPlan* plan, NodeVector inputs, std::vector<std::string> input_labels,
+ const AsofJoinNodeOptions& join_options,
+ std::shared_ptr<Schema> output_schema,
+ std::unique_ptr<AsofJoinSchema> schema_mgr);
+
+ virtual ~AsofJoinNode() {
+ _process.push(false); // poison pill
+ _process_thread.join();
+ }
+
+ static arrow::Result<ExecNode*> Make(ExecPlan* plan, std::vector<ExecNode*> inputs,
+ const ExecNodeOptions& options) {
+ std::unique_ptr<AsofJoinSchema> schema_mgr =
+ ::arrow::internal::make_unique<AsofJoinSchema>();
Review Comment:
This looks like it's only used to create the output schema so I'm not sure why you are storing it as a member variable. Could this be a static method? Or, could you just create a local variable here, generate the output schema, and then let it be discarded?
##########
cpp/src/arrow/compute/exec/asof_join_node.cc:
##########
@@ -0,0 +1,800 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+#include <iostream>
+#include <unordered_map>
+
+#include <arrow/api.h>
+#include <arrow/compute/api.h>
+#include <arrow/util/optional.h>
+#include "arrow/compute/exec/asof_join.h"
+#include "arrow/compute/exec/exec_plan.h"
+#include "arrow/compute/exec/options.h"
+#include "arrow/compute/exec/schema_util.h"
+#include "arrow/compute/exec/util.h"
+#include "arrow/result.h"
+#include "arrow/status.h"
+#include "arrow/util/checked_cast.h"
+#include "arrow/util/future.h"
+#include "arrow/util/make_unique.h"
+
+#include <condition_variable>
+#include <mutex>
+#include <thread>
+
+namespace arrow {
+namespace compute {
+
+/**
+ * Simple implementation for an unbound concurrent queue
+ */
+template <class T>
+class ConcurrentQueue {
+ public:
+ T pop() {
+ std::unique_lock<std::mutex> lock(mutex_);
+ cond_.wait(lock, [&] { return !queue_.empty(); });
+ auto item = queue_.front();
+ queue_.pop();
+ return item;
+ }
+
+ void push(const T& item) {
+ std::unique_lock<std::mutex> lock(mutex_);
+ queue_.push(item);
+ cond_.notify_one();
+ }
+
+ util::optional<T> try_pop() {
+ // Try to pop the oldest value from the queue (or return nullopt if none)
+ std::unique_lock<std::mutex> lock(mutex_);
+ if (queue_.empty()) {
+ return util::nullopt;
+ } else {
+ auto item = queue_.front();
+ queue_.pop();
+ return item;
+ }
+ }
+
+ bool empty() const {
+ std::unique_lock<std::mutex> lock(mutex_);
+ return queue_.empty();
+ }
+
+ // Un-synchronized access to front
+ // For this to be "safe":
+ // 1) the caller logically guarantees that queue is not empty
+ // 2) pop/try_pop cannot be called concurrently with this
+ const T& unsync_front() const { return queue_.front(); }
+
+ private:
+ std::queue<T> queue_;
+ mutable std::mutex mutex_;
+ std::condition_variable cond_;
+};
+
+struct MemoStore {
+ // Stores last known values for all the keys
+
+ struct Entry {
+ // Timestamp associated with the entry
+ int64_t _time;
+
+ // Batch associated with the entry (perf is probably OK for this; batches change
+ // rarely)
+ std::shared_ptr<arrow::RecordBatch> _batch;
+
+ // Row associated with the entry
+ row_index_t _row;
+ };
+
+ std::unordered_map<KeyType, Entry> _entries;
+
+ void store(const std::shared_ptr<RecordBatch>& batch, row_index_t row, int64_t time,
+ KeyType key) {
+ auto& e = _entries[key];
+ // that we can do this assignment optionally, is why we
+ // can get array 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;
+ }
+
+ util::optional<const Entry*> get_entry_for_key(KeyType key) const {
+ auto e = _entries.find(key);
+ if (_entries.end() == e) return util::nullopt;
+ return util::optional<const Entry*>(&e->second);
+ }
+
+ void remove_entries_with_lesser_time(int64_t ts) {
+ size_t dbg_size0 = _entries.size();
+ for (auto e = _entries.begin(); e != _entries.end();)
+ if (e->second._time < ts)
+ e = _entries.erase(e);
+ else
+ ++e;
+ size_t dbg_size1 = _entries.size();
+ if (dbg_size1 < dbg_size0) {
+ // cerr << "Removed " << dbg_size0-dbg_size1 << " memo entries.\n";
+ }
+ }
+};
+
+class InputState {
+ // InputState correponds to an input
+ // Input record batches are queued up in InputState until processed and
+ // turned into output record batches.
+
+ public:
+ InputState(const std::shared_ptr<arrow::Schema>& schema,
+ const std::string& time_col_name, const std::string& key_col_name,
+ util::optional<KeyType> wildcard_key)
+ : queue_(),
+ wildcard_key_(wildcard_key),
+ schema_(schema),
+ time_col_index_(
+ schema->GetFieldIndex(time_col_name)), // TODO: handle missing field name
+ key_col_index_(schema->GetFieldIndex(key_col_name)) {}
+
+ col_index_t init_src_to_dst_mapping(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 && is_time_or_key_column(i)))
+ src_to_dst_[i] = dst_offset++;
+ return dst_offset;
+ }
+
+ const util::optional<col_index_t>& map_src_to_dst(col_index_t src) const {
+ return src_to_dst_[src];
+ }
+
+ bool is_time_or_key_column(col_index_t i) const {
+ assert(i < schema_->num_fields());
+ return (i == time_col_index_) || (i == key_col_index_);
+ }
+
+ // Gets the latest row index, assuming the queue isn't empty
+ row_index_t get_latest_row() const { return latest_ref_row_; }
+
+ bool empty() const {
+ if (latest_ref_row_ > 0)
+ return false; // cannot be empty if ref row is >0 -- can avoid slow queue lock
+ // below
+ return queue_.empty();
+ }
+
+ int countbatches_processed_() const { return batches_processed_; }
+ int count_total_batches() const { return total_batches_; }
+
+ // Gets latest batch (precondition: must not be empty)
+ const std::shared_ptr<arrow::RecordBatch>& get_latest_batch() const {
+ return queue_.unsync_front();
+ }
+ KeyType get_latest_key() const {
+ return queue_.unsync_front()
+ ->column_data(key_col_index_)
+ ->GetValues<KeyType>(1)[latest_ref_row_];
+ }
+ int64_t get_latest_time() const {
+ return queue_.unsync_front()
+ ->column_data(time_col_index_)
+ ->GetValues<int64_t>(1)[latest_ref_row_];
+ }
+
+ bool finished() const { return batches_processed_ == total_batches_; }
+
+ bool advance() {
+ // 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();
+ if (have_active_batch) {
+ // If we have an active batch
+ if (++latest_ref_row_ >= (row_index_t)queue_.unsync_front()->num_rows()) {
+ // hit the end of the batch, need to get the next batch if possible.
+ ++batches_processed_;
+ latest_ref_row_ = 0;
+ have_active_batch &= !queue_.try_pop();
+ if (have_active_batch)
+ assert(queue_.unsync_front()->num_rows() > 0); // empty batches disallowed
+ }
+ }
+ return have_active_batch;
+ }
+
+ // Advance the data to be immediately past the specified TS, updating latest and
+ // latest_ref_row to the latest data prior to that immediate just past Returns true if
+ // updates were made, false if not.
+ bool advance_and_memoize(int64_t ts) {
+ // Advance the right side row index until we reach the latest right row (for each key)
+ // for the given left timestamp.
+
+ // Check if already updated for TS (or if there is no latest)
+ if (empty()) return false; // can't advance if empty
+ auto latest_time = get_latest_time();
+ if (latest_time > ts) return false; // already advanced
+
+ // Not updated. Try to update and possibly advance.
+ bool updated = false;
+ do {
+ latest_time = get_latest_time();
+ // 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 <= ts) {
+ memo_.store(get_latest_batch(), latest_ref_row_, latest_time, get_latest_key());
+ } else {
+ break; // hit a future timestamp -- done updating for now
+ }
+ updated = true;
+ } while (advance());
+ return updated;
+ }
+
+ void push(const std::shared_ptr<arrow::RecordBatch>& rb) {
+ if (rb->num_rows() > 0) {
+ queue_.push(rb);
+ } else {
+ ++batches_processed_; // don't enqueue empty batches, just record as processed
+ }
+ }
+
+ util::optional<const MemoStore::Entry*> get_memo_entry_for_key(KeyType key) {
+ auto r = memo_.get_entry_for_key(key);
+ if (r.has_value()) return r;
+ if (wildcard_key_.has_value()) r = memo_.get_entry_for_key(*wildcard_key_);
+ return r;
+ }
+
+ util::optional<int64_t> get_memo_time_for_key(KeyType key) {
+ auto r = get_memo_entry_for_key(key);
+ return r.has_value() ? util::make_optional((*r)->_time) : util::nullopt;
+ }
+
+ void remove_memo_entries_with_lesser_time(int64_t ts) {
+ memo_.remove_entries_with_lesser_time(ts);
+ }
+
+ const std::shared_ptr<Schema>& get_schema() const { return schema_; }
+
+ void set_total_batches(int n) {
+ assert(n >= 0);
+ assert(total_batches_ == -1); // shouldn't be set more than once
+ total_batches_ = n;
+ }
+
+ private:
+ // Pending record batches. The latest is the front. Batches cannot be empty.
+ ConcurrentQueue<std::shared_ptr<RecordBatch>> queue_;
+
+ // Wildcard key for this input, if applicable.
+ util::optional<KeyType> wildcard_key_;
+
+ // Schema associated with the input
+ std::shared_ptr<Schema> schema_;
+
+ // Total number of batches (only int because InputFinished uses int)
+ int total_batches_ = -1;
+
+ // Number of batches processed so far (only int because InputFinished uses int)
+ int batches_processed_ = 0;
+
+ // Index of the time col
+ col_index_t time_col_index_;
+
+ // Index of the key col
+ col_index_t key_col_index_;
+
+ // Index of the latest row reference within; if >0 then queue_ cannot be empty
+ row_index_t latest_ref_row_ =
+ 0; // must be < queue_.front()->num_rows() if queue_ is non-empty
+
+ // Stores latest known values for the various keys
+ MemoStore memo_;
+
+ // Mapping of source columns to destination columns
+ std::vector<util::optional<col_index_t>> src_to_dst_;
+};
+
+template <size_t MAX_TABLES>
+struct CompositeReferenceRow {
+ struct Entry {
+ arrow::RecordBatch* batch; // can be NULL if there's no value
+ row_index_t row;
+ };
+ Entry refs[MAX_TABLES];
+};
+
+// A table of composite reference rows. Rows maintain pointers to the
+// constituent record batches, but the overall table retains shared_ptr
+// references to ensure memory remains resident while the table is live.
+//
+// The main reason for this is that, especially for wide tables, joins
+// are effectively row-oriented, rather than column-oriented. Separating
+// the join part from the columnar materialization part simplifies the
+// logic around data types and increases efficiency.
+//
+// We don't put the shared_ptr's into the rows for efficiency reasons.
+template <size_t MAX_TABLES>
+class CompositeReferenceTable {
+ public:
+ explicit CompositeReferenceTable(size_t n_tables) : n_tables_(n_tables) {
+ assert(n_tables_ >= 1);
+ assert(n_tables_ <= MAX_TABLES);
+ }
+
+ size_t n_rows() const { return rows_.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<std::unique_ptr<InputState>>& in, int64_t tolerance) {
+ assert(in.size() == n_tables_);
+
+ // Get the LHS key
+ KeyType key = in[0]->get_latest_key();
+
+ // Add row and setup LHS
+ // (the LHS state comes just from the latest row of the LHS table)
+ assert(!in[0]->empty());
+ const std::shared_ptr<arrow::RecordBatch>& lhs_latest_batch =
+ in[0]->get_latest_batch();
+ row_index_t lhs_latest_row = in[0]->get_latest_row();
+ int64_t lhs_latest_time = in[0]->get_latest_time();
+ 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 = rows_.size() + new_batch_size;
+ if (rows_.capacity() < new_capacity) rows_.reserve(new_capacity);
+ }
+ rows_.resize(rows_.size() + 1);
+ auto& row = rows_.back();
+ row.refs[0].batch = lhs_latest_batch.get();
+ row.refs[0].row = lhs_latest_row;
+ add_record_batch_ref(lhs_latest_batch);
+
+ // 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) {
+ util::optional<const MemoStore::Entry*> opt_entry =
+ in[i]->get_memo_entry_for_key(key);
+ if (opt_entry.has_value()) {
+ assert(*opt_entry);
+ if ((*opt_entry)->_time + tolerance >= lhs_latest_time) {
+ // Have a valid entry
+ const MemoStore::Entry* entry = *opt_entry;
+ row.refs[i].batch = entry->_batch.get();
+ row.refs[i].row = entry->_row;
+ add_record_batch_ref(entry->_batch);
+ continue;
+ }
+ }
+ row.refs[i].batch = NULL;
+ row.refs[i].row = 0;
+ }
+ }
+
+ // Materializes the current reference table into a target record batch
+ Result<std::shared_ptr<RecordBatch>> materialize(
+ const std::shared_ptr<arrow::Schema>& output_schema,
+ const std::vector<std::unique_ptr<InputState>>& state) {
+ // cerr << "materialize BEGIN\n";
+ assert(state.size() == n_tables_);
+ assert(state.size() >= 1);
+
+ // Don't build empty batches
+ size_t n_rows = rows_.size();
+ if (!n_rows) return NULLPTR;
+
+ // Build the arrays column-by-column from the rows
+ std::vector<std::shared_ptr<arrow::Array>> arrays(output_schema->num_fields());
+ for (size_t i_table = 0; i_table < n_tables_; ++i_table) {
+ int n_src_cols = state.at(i_table)->get_schema()->num_fields();
+ {
+ for (col_index_t i_src_col = 0; i_src_col < n_src_cols; ++i_src_col) {
+ util::optional<col_index_t> i_dst_col_opt =
+ state[i_table]->map_src_to_dst(i_src_col);
+ if (!i_dst_col_opt) continue;
+ col_index_t i_dst_col = *i_dst_col_opt;
+ const auto& src_field = state[i_table]->get_schema()->field(i_src_col);
+ const auto& dst_field = output_schema->field(i_dst_col);
+ assert(src_field->type()->Equals(dst_field->type()));
+ assert(src_field->name() == dst_field->name());
+ const auto& field_type = src_field->type();
+
+ if (field_type->Equals(arrow::int32())) {
+ ARROW_ASSIGN_OR_RAISE(
+ arrays.at(i_dst_col),
+ (materialize_primitive_column<arrow::Int32Builder, int32_t>(i_table,
+ i_src_col)));
+ } else if (field_type->Equals(arrow::int64())) {
+ ARROW_ASSIGN_OR_RAISE(
+ arrays.at(i_dst_col),
+ (materialize_primitive_column<arrow::Int64Builder, int64_t>(i_table,
+ i_src_col)));
+ } else if (field_type->Equals(arrow::float64())) {
+ ARROW_ASSIGN_OR_RAISE(
+ arrays.at(i_dst_col),
+ (materialize_primitive_column<arrow::DoubleBuilder, double>(i_table,
+ i_src_col)));
+ } else {
+ ARROW_RETURN_NOT_OK(
+ Status::Invalid("Unsupported data type: ", src_field->name()));
+ }
+ }
+ }
+ }
+
+ // Build the result
+ assert(sizeof(size_t) >= sizeof(int64_t)); // Make takes signed int64_t for num_rows
+
+ // TODO: check n_rows for cast
+ std::shared_ptr<arrow::RecordBatch> r =
+ arrow::RecordBatch::Make(output_schema, (int64_t)n_rows, arrays);
+ return r;
+ }
+
+ // Returns true if there are no rows
+ bool empty() const { return rows_.empty(); }
+
+ private:
+ // Contains shared_ptr refs for all RecordBatches referred to by the contents of rows_
+ std::unordered_map<uintptr_t, std::shared_ptr<RecordBatch>> _ptr2ref;
+
+ // Row table references
+ std::vector<CompositeReferenceRow<MAX_TABLES>> rows_;
+
+ // Total number of tables in the composite table
+ size_t n_tables_;
+
+ // Adds a RecordBatch ref to the mapping, if needed
+ void add_record_batch_ref(const std::shared_ptr<RecordBatch>& ref) {
+ if (!_ptr2ref.count((uintptr_t)ref.get())) _ptr2ref[(uintptr_t)ref.get()] = ref;
+ }
+
+ template <class Builder, class PrimitiveType>
+ Result<std::shared_ptr<Array>> materialize_primitive_column(size_t i_table,
+ col_index_t i_col) {
+ Builder builder;
+ ARROW_RETURN_NOT_OK(builder.Reserve(rows_.size()));
+ for (row_index_t i_row = 0; i_row < rows_.size(); ++i_row) {
+ const auto& ref = rows_[i_row].refs[i_table];
+ if (ref.batch) {
+ builder.UnsafeAppend(
+ ref.batch->column_data(i_col)->template GetValues<PrimitiveType>(1)[ref.row]);
+ } else {
+ builder.UnsafeAppendNull();
+ }
+ }
+ std::shared_ptr<Array> result;
+ ARROW_RETURN_NOT_OK(builder.Finish(&result));
+ return result;
+ }
+};
+
+class AsofJoinNode : public ExecNode {
+ // Constructs labels for inputs
+ static std::vector<std::string> build_input_labels(
+ const std::vector<ExecNode*>& inputs) {
+ std::vector<std::string> r(inputs.size());
+ for (size_t i = 0; i < r.size(); ++i) r[i] = "input_" + std::to_string(i) + "_label";
+ return r;
+ }
+
+ // Advances the RHS as far as possible to be up to date for the current LHS timestamp
+ bool update_rhs() {
+ auto& lhs = *_state.at(0);
+ auto lhs_latest_time = lhs.get_latest_time();
+ bool any_updated = false;
+ for (size_t i = 1; i < _state.size(); ++i)
+ any_updated |= _state[i]->advance_and_memoize(lhs_latest_time);
+ return any_updated;
+ }
+
+ // Returns false if RHS not up to date for LHS
+ bool is_up_to_date_for_lhs_row() const {
+ auto& lhs = *_state[0];
+ if (lhs.empty()) return false; // can't proceed if nothing on the LHS
+ int64_t lhs_ts = lhs.get_latest_time();
+ for (size_t i = 1; i < _state.size(); ++i) {
+ auto& rhs = *_state[i];
+ if (!rhs.finished()) {
+ // If RHS is finished, then we know it's up to date (but if it isn't, it might be
+ // up to date)
+ if (rhs.empty())
+ return false; // RHS isn't finished, but is empty --> not up to date
+ if (lhs_ts >= rhs.get_latest_time())
+ return false; // TS not up to date (and not finished)
+ }
+ }
+ return true;
+ }
+
+ Result<std::shared_ptr<RecordBatch>> process_inner() {
+ assert(!_state.empty());
+ auto& lhs = *_state.at(0);
+
+ // Construct new target table if needed
+ CompositeReferenceTable<MAX_JOIN_TABLES> dst(_state.size());
+
+ // 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;
+
+ // Advance each of the RHS as far as possible to be up to date for the LHS timestamp
+ bool any_advanced = update_rhs();
+
+ // Only update if we have up-to-date information for the LHS row
+ if (is_up_to_date_for_lhs_row()) {
+ dst.emplace(_state, _options.tolerance);
+ if (!lhs.advance()) break; // if we can't advance LHS, we're done for this batch
+ } else {
+ if ((!any_advanced) && (_state.size() > 1)) break; // need to wait for new data
+ }
+ }
+
+ // Prune memo entries that have expired (to bound memory consumption)
+ if (!lhs.empty()) {
+ for (size_t i = 1; i < _state.size(); ++i) {
+ _state[i]->remove_memo_entries_with_lesser_time(lhs.get_latest_time() -
+ _options.tolerance);
+ }
+ }
+
+ // Emit the batch
+ if (dst.empty()) {
+ return NULLPTR;
+ } else {
+ return dst.materialize(output_schema(), _state);
+ }
+ }
+
+ void process() {
+ std::cerr << "process() begin\n";
+
+ std::lock_guard<std::mutex> guard(_gate);
+ if (finished_.is_finished()) {
+ std::cerr << "InputReceived EARLYEND\n";
+ return;
+ }
+
+ // Process batches while we have data
+ for (;;) {
+ Result<std::shared_ptr<RecordBatch>> result = process_inner();
+
+ if (result.ok()) {
+ auto out_rb = *result;
+ if (!out_rb) break;
+ ++_progress_batches_produced;
+ ExecBatch out_b(*out_rb);
+ outputs_[0]->InputReceived(this, std::move(out_b));
+ } else {
+ StopProducing();
+ ErrorIfNotOk(result.status());
+ return;
+ }
+ }
+
+ std::cerr << "process() end\n";
+
+ // 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 (_state.at(0)->finished()) {
+ total_batches_produced_ = util::make_optional<int>(_progress_batches_produced);
+ StopProducing();
+ assert(total_batches_produced_.has_value());
+ outputs_[0]->InputFinished(this, *total_batches_produced_);
+ }
+ }
+
+ void process_thread() {
+ std::cerr << "AsofJoinNode::process_thread started.\n";
+ for (;;) {
+ if (!_process.pop()) {
+ std::cerr << "AsofJoinNode::process_thread done.\n";
+ return;
+ }
+ process();
+ }
+ }
+
+ static void process_thread_wrapper(AsofJoinNode* node) { node->process_thread(); }
+
+ public:
+ AsofJoinNode(ExecPlan* plan, NodeVector inputs, std::vector<std::string> input_labels,
+ const AsofJoinNodeOptions& join_options,
+ std::shared_ptr<Schema> output_schema,
+ std::unique_ptr<AsofJoinSchema> schema_mgr);
+
+ virtual ~AsofJoinNode() {
+ _process.push(false); // poison pill
+ _process_thread.join();
+ }
+
+ static arrow::Result<ExecNode*> Make(ExecPlan* plan, std::vector<ExecNode*> inputs,
+ const ExecNodeOptions& options) {
+ std::unique_ptr<AsofJoinSchema> schema_mgr =
+ ::arrow::internal::make_unique<AsofJoinSchema>();
+
+ const auto& join_options = checked_cast<const AsofJoinNodeOptions&>(options);
+ std::shared_ptr<Schema> output_schema =
+ schema_mgr->MakeOutputSchema(inputs, join_options);
+
+ std::vector<std::string> input_labels(inputs.size());
+ input_labels[0] = "left";
+ for (size_t i = 1; i < inputs.size(); ++i) {
+ input_labels[i] = "right_" + std::to_string(i);
+ }
+
+ return plan->EmplaceNode<AsofJoinNode>(plan, inputs, std::move(input_labels),
+ join_options, std::move(output_schema),
+ std::move(schema_mgr));
+ }
+
+ const char* kind_name() const override { return "AsofJoinNode"; }
+
+ void InputReceived(ExecNode* input, ExecBatch batch) override {
+ // Get the input
+ ARROW_DCHECK(std::find(inputs_.begin(), inputs_.end(), input) != inputs_.end());
+ size_t k = std::find(inputs_.begin(), inputs_.end(), input) - inputs_.begin();
+ std::cerr << "InputReceived BEGIN (k=" << k << ")\n";
+
+ // Put into the queue
+ auto rb = *batch.ToRecordBatch(input->output_schema());
+
+ _state.at(k)->push(rb);
+ _process.push(true);
+
+ std::cerr << "InputReceived END\n";
+ }
+ void ErrorReceived(ExecNode* input, Status error) override {
+ outputs_[0]->ErrorReceived(this, std::move(error));
+ StopProducing();
+ }
+ void InputFinished(ExecNode* input, int total_batches) override {
+ std::cerr << "InputFinished BEGIN\n";
+ // bool is_finished=false;
+ {
+ std::lock_guard<std::mutex> guard(_gate);
+ std::cerr << "InputFinished find\n";
+ ARROW_DCHECK(std::find(inputs_.begin(), inputs_.end(), input) != inputs_.end());
+ size_t k = std::find(inputs_.begin(), inputs_.end(), input) - inputs_.begin();
+ // cerr << "set_total_batches for input " << k << ": " << total_batches << "\n";
+ _state.at(k)->set_total_batches(total_batches);
+ }
+ // 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.
+ _process.push(true);
+
+ std::cerr << "InputFinished END\n";
+ }
+ Status StartProducing() override {
+ std::cout << "StartProducing"
+ << "\n";
+ finished_ = arrow::Future<>::Make();
+ return Status::OK();
+ }
+ void PauseProducing(ExecNode* output, int32_t counter) override {
+ std::cout << "PauseProducing"
+ << "\n";
+ }
+ void ResumeProducing(ExecNode* output, int32_t counter) override {
+ std::cout << "ResumeProducing"
+ << "\n";
+ }
+ void StopProducing(ExecNode* output) override {
+ DCHECK_EQ(output, outputs_[0]);
+ StopProducing();
+ std::cout << "StopProducing"
+ << "\n";
+ }
+ void StopProducing() override {
+ std::cerr << "StopProducing" << std::endl;
+ // if(batch_count_.Cancel()) finished_.MarkFinished();
+ finished_.MarkFinished();
+ for (auto&& input : inputs_) input->StopProducing(this);
Review Comment:
```suggestion
```
##########
cpp/src/arrow/compute/exec/asof_join_node.cc:
##########
@@ -0,0 +1,800 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+#include <iostream>
+#include <unordered_map>
+
+#include <arrow/api.h>
+#include <arrow/compute/api.h>
+#include <arrow/util/optional.h>
+#include "arrow/compute/exec/asof_join.h"
+#include "arrow/compute/exec/exec_plan.h"
+#include "arrow/compute/exec/options.h"
+#include "arrow/compute/exec/schema_util.h"
+#include "arrow/compute/exec/util.h"
+#include "arrow/result.h"
+#include "arrow/status.h"
+#include "arrow/util/checked_cast.h"
+#include "arrow/util/future.h"
+#include "arrow/util/make_unique.h"
+
+#include <condition_variable>
+#include <mutex>
+#include <thread>
+
+namespace arrow {
+namespace compute {
+
+/**
+ * Simple implementation for an unbound concurrent queue
+ */
+template <class T>
+class ConcurrentQueue {
+ public:
+ T pop() {
+ std::unique_lock<std::mutex> lock(mutex_);
+ cond_.wait(lock, [&] { return !queue_.empty(); });
+ auto item = queue_.front();
+ queue_.pop();
+ return item;
+ }
+
+ void push(const T& item) {
+ std::unique_lock<std::mutex> lock(mutex_);
+ queue_.push(item);
+ cond_.notify_one();
+ }
+
+ util::optional<T> try_pop() {
+ // Try to pop the oldest value from the queue (or return nullopt if none)
+ std::unique_lock<std::mutex> lock(mutex_);
+ if (queue_.empty()) {
+ return util::nullopt;
+ } else {
+ auto item = queue_.front();
+ queue_.pop();
+ return item;
+ }
+ }
+
+ bool empty() const {
+ std::unique_lock<std::mutex> lock(mutex_);
+ return queue_.empty();
+ }
+
+ // Un-synchronized access to front
+ // For this to be "safe":
+ // 1) the caller logically guarantees that queue is not empty
+ // 2) pop/try_pop cannot be called concurrently with this
+ const T& unsync_front() const { return queue_.front(); }
+
+ private:
+ std::queue<T> queue_;
+ mutable std::mutex mutex_;
+ std::condition_variable cond_;
+};
+
+struct MemoStore {
+ // Stores last known values for all the keys
+
+ struct Entry {
+ // Timestamp associated with the entry
+ int64_t _time;
+
+ // Batch associated with the entry (perf is probably OK for this; batches change
+ // rarely)
+ std::shared_ptr<arrow::RecordBatch> _batch;
+
+ // Row associated with the entry
+ row_index_t _row;
+ };
+
+ std::unordered_map<KeyType, Entry> _entries;
+
+ void store(const std::shared_ptr<RecordBatch>& batch, row_index_t row, int64_t time,
+ KeyType key) {
+ auto& e = _entries[key];
+ // that we can do this assignment optionally, is why we
+ // can get array 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;
+ }
+
+ util::optional<const Entry*> get_entry_for_key(KeyType key) const {
+ auto e = _entries.find(key);
+ if (_entries.end() == e) return util::nullopt;
+ return util::optional<const Entry*>(&e->second);
+ }
+
+ void remove_entries_with_lesser_time(int64_t ts) {
+ size_t dbg_size0 = _entries.size();
+ for (auto e = _entries.begin(); e != _entries.end();)
+ if (e->second._time < ts)
+ e = _entries.erase(e);
+ else
+ ++e;
+ size_t dbg_size1 = _entries.size();
+ if (dbg_size1 < dbg_size0) {
+ // cerr << "Removed " << dbg_size0-dbg_size1 << " memo entries.\n";
+ }
+ }
+};
+
+class InputState {
+ // InputState correponds to an input
+ // Input record batches are queued up in InputState until processed and
+ // turned into output record batches.
+
+ public:
+ InputState(const std::shared_ptr<arrow::Schema>& schema,
+ const std::string& time_col_name, const std::string& key_col_name,
+ util::optional<KeyType> wildcard_key)
+ : queue_(),
+ wildcard_key_(wildcard_key),
+ schema_(schema),
+ time_col_index_(
+ schema->GetFieldIndex(time_col_name)), // TODO: handle missing field name
+ key_col_index_(schema->GetFieldIndex(key_col_name)) {}
+
+ col_index_t init_src_to_dst_mapping(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 && is_time_or_key_column(i)))
+ src_to_dst_[i] = dst_offset++;
+ return dst_offset;
+ }
+
+ const util::optional<col_index_t>& map_src_to_dst(col_index_t src) const {
+ return src_to_dst_[src];
+ }
+
+ bool is_time_or_key_column(col_index_t i) const {
+ assert(i < schema_->num_fields());
+ return (i == time_col_index_) || (i == key_col_index_);
+ }
+
+ // Gets the latest row index, assuming the queue isn't empty
+ row_index_t get_latest_row() const { return latest_ref_row_; }
+
+ bool empty() const {
+ if (latest_ref_row_ > 0)
+ return false; // cannot be empty if ref row is >0 -- can avoid slow queue lock
+ // below
+ return queue_.empty();
+ }
+
+ int countbatches_processed_() const { return batches_processed_; }
+ int count_total_batches() const { return total_batches_; }
+
+ // Gets latest batch (precondition: must not be empty)
+ const std::shared_ptr<arrow::RecordBatch>& get_latest_batch() const {
+ return queue_.unsync_front();
+ }
+ KeyType get_latest_key() const {
+ return queue_.unsync_front()
+ ->column_data(key_col_index_)
+ ->GetValues<KeyType>(1)[latest_ref_row_];
+ }
+ int64_t get_latest_time() const {
+ return queue_.unsync_front()
+ ->column_data(time_col_index_)
+ ->GetValues<int64_t>(1)[latest_ref_row_];
+ }
+
+ bool finished() const { return batches_processed_ == total_batches_; }
+
+ bool advance() {
+ // 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();
+ if (have_active_batch) {
+ // If we have an active batch
+ if (++latest_ref_row_ >= (row_index_t)queue_.unsync_front()->num_rows()) {
+ // hit the end of the batch, need to get the next batch if possible.
+ ++batches_processed_;
+ latest_ref_row_ = 0;
+ have_active_batch &= !queue_.try_pop();
+ if (have_active_batch)
+ assert(queue_.unsync_front()->num_rows() > 0); // empty batches disallowed
+ }
+ }
+ return have_active_batch;
+ }
+
+ // Advance the data to be immediately past the specified TS, updating latest and
+ // latest_ref_row to the latest data prior to that immediate just past Returns true if
+ // updates were made, false if not.
+ bool advance_and_memoize(int64_t ts) {
+ // Advance the right side row index until we reach the latest right row (for each key)
+ // for the given left timestamp.
+
+ // Check if already updated for TS (or if there is no latest)
+ if (empty()) return false; // can't advance if empty
+ auto latest_time = get_latest_time();
+ if (latest_time > ts) return false; // already advanced
+
+ // Not updated. Try to update and possibly advance.
+ bool updated = false;
+ do {
+ latest_time = get_latest_time();
+ // 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 <= ts) {
+ memo_.store(get_latest_batch(), latest_ref_row_, latest_time, get_latest_key());
+ } else {
+ break; // hit a future timestamp -- done updating for now
+ }
+ updated = true;
+ } while (advance());
+ return updated;
+ }
+
+ void push(const std::shared_ptr<arrow::RecordBatch>& rb) {
+ if (rb->num_rows() > 0) {
+ queue_.push(rb);
+ } else {
+ ++batches_processed_; // don't enqueue empty batches, just record as processed
+ }
+ }
+
+ util::optional<const MemoStore::Entry*> get_memo_entry_for_key(KeyType key) {
+ auto r = memo_.get_entry_for_key(key);
+ if (r.has_value()) return r;
+ if (wildcard_key_.has_value()) r = memo_.get_entry_for_key(*wildcard_key_);
+ return r;
+ }
+
+ util::optional<int64_t> get_memo_time_for_key(KeyType key) {
+ auto r = get_memo_entry_for_key(key);
+ return r.has_value() ? util::make_optional((*r)->_time) : util::nullopt;
+ }
+
+ void remove_memo_entries_with_lesser_time(int64_t ts) {
+ memo_.remove_entries_with_lesser_time(ts);
+ }
+
+ const std::shared_ptr<Schema>& get_schema() const { return schema_; }
+
+ void set_total_batches(int n) {
+ assert(n >= 0);
+ assert(total_batches_ == -1); // shouldn't be set more than once
+ total_batches_ = n;
+ }
+
+ private:
+ // Pending record batches. The latest is the front. Batches cannot be empty.
+ ConcurrentQueue<std::shared_ptr<RecordBatch>> queue_;
+
+ // Wildcard key for this input, if applicable.
+ util::optional<KeyType> wildcard_key_;
+
+ // Schema associated with the input
+ std::shared_ptr<Schema> schema_;
+
+ // Total number of batches (only int because InputFinished uses int)
+ int total_batches_ = -1;
+
+ // Number of batches processed so far (only int because InputFinished uses int)
+ int batches_processed_ = 0;
+
+ // Index of the time col
+ col_index_t time_col_index_;
+
+ // Index of the key col
+ col_index_t key_col_index_;
+
+ // Index of the latest row reference within; if >0 then queue_ cannot be empty
+ row_index_t latest_ref_row_ =
+ 0; // must be < queue_.front()->num_rows() if queue_ is non-empty
+
+ // Stores latest known values for the various keys
+ MemoStore memo_;
+
+ // Mapping of source columns to destination columns
+ std::vector<util::optional<col_index_t>> src_to_dst_;
+};
+
+template <size_t MAX_TABLES>
+struct CompositeReferenceRow {
+ struct Entry {
+ arrow::RecordBatch* batch; // can be NULL if there's no value
+ row_index_t row;
+ };
+ Entry refs[MAX_TABLES];
+};
+
+// A table of composite reference rows. Rows maintain pointers to the
+// constituent record batches, but the overall table retains shared_ptr
+// references to ensure memory remains resident while the table is live.
+//
+// The main reason for this is that, especially for wide tables, joins
+// are effectively row-oriented, rather than column-oriented. Separating
+// the join part from the columnar materialization part simplifies the
+// logic around data types and increases efficiency.
+//
+// We don't put the shared_ptr's into the rows for efficiency reasons.
+template <size_t MAX_TABLES>
+class CompositeReferenceTable {
+ public:
+ explicit CompositeReferenceTable(size_t n_tables) : n_tables_(n_tables) {
+ assert(n_tables_ >= 1);
+ assert(n_tables_ <= MAX_TABLES);
+ }
+
+ size_t n_rows() const { return rows_.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<std::unique_ptr<InputState>>& in, int64_t tolerance) {
+ assert(in.size() == n_tables_);
+
+ // Get the LHS key
+ KeyType key = in[0]->get_latest_key();
+
+ // Add row and setup LHS
+ // (the LHS state comes just from the latest row of the LHS table)
+ assert(!in[0]->empty());
+ const std::shared_ptr<arrow::RecordBatch>& lhs_latest_batch =
+ in[0]->get_latest_batch();
+ row_index_t lhs_latest_row = in[0]->get_latest_row();
+ int64_t lhs_latest_time = in[0]->get_latest_time();
+ 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 = rows_.size() + new_batch_size;
+ if (rows_.capacity() < new_capacity) rows_.reserve(new_capacity);
+ }
+ rows_.resize(rows_.size() + 1);
+ auto& row = rows_.back();
+ row.refs[0].batch = lhs_latest_batch.get();
+ row.refs[0].row = lhs_latest_row;
+ add_record_batch_ref(lhs_latest_batch);
+
+ // 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) {
+ util::optional<const MemoStore::Entry*> opt_entry =
+ in[i]->get_memo_entry_for_key(key);
+ if (opt_entry.has_value()) {
+ assert(*opt_entry);
+ if ((*opt_entry)->_time + tolerance >= lhs_latest_time) {
+ // Have a valid entry
+ const MemoStore::Entry* entry = *opt_entry;
+ row.refs[i].batch = entry->_batch.get();
+ row.refs[i].row = entry->_row;
+ add_record_batch_ref(entry->_batch);
+ continue;
+ }
+ }
+ row.refs[i].batch = NULL;
+ row.refs[i].row = 0;
+ }
+ }
+
+ // Materializes the current reference table into a target record batch
+ Result<std::shared_ptr<RecordBatch>> materialize(
+ const std::shared_ptr<arrow::Schema>& output_schema,
+ const std::vector<std::unique_ptr<InputState>>& state) {
+ // cerr << "materialize BEGIN\n";
+ assert(state.size() == n_tables_);
+ assert(state.size() >= 1);
+
+ // Don't build empty batches
+ size_t n_rows = rows_.size();
+ if (!n_rows) return NULLPTR;
+
+ // Build the arrays column-by-column from the rows
+ std::vector<std::shared_ptr<arrow::Array>> arrays(output_schema->num_fields());
+ for (size_t i_table = 0; i_table < n_tables_; ++i_table) {
+ int n_src_cols = state.at(i_table)->get_schema()->num_fields();
+ {
+ for (col_index_t i_src_col = 0; i_src_col < n_src_cols; ++i_src_col) {
+ util::optional<col_index_t> i_dst_col_opt =
+ state[i_table]->map_src_to_dst(i_src_col);
+ if (!i_dst_col_opt) continue;
+ col_index_t i_dst_col = *i_dst_col_opt;
+ const auto& src_field = state[i_table]->get_schema()->field(i_src_col);
+ const auto& dst_field = output_schema->field(i_dst_col);
+ assert(src_field->type()->Equals(dst_field->type()));
+ assert(src_field->name() == dst_field->name());
+ const auto& field_type = src_field->type();
+
+ if (field_type->Equals(arrow::int32())) {
+ ARROW_ASSIGN_OR_RAISE(
+ arrays.at(i_dst_col),
+ (materialize_primitive_column<arrow::Int32Builder, int32_t>(i_table,
+ i_src_col)));
+ } else if (field_type->Equals(arrow::int64())) {
+ ARROW_ASSIGN_OR_RAISE(
+ arrays.at(i_dst_col),
+ (materialize_primitive_column<arrow::Int64Builder, int64_t>(i_table,
+ i_src_col)));
+ } else if (field_type->Equals(arrow::float64())) {
+ ARROW_ASSIGN_OR_RAISE(
+ arrays.at(i_dst_col),
+ (materialize_primitive_column<arrow::DoubleBuilder, double>(i_table,
+ i_src_col)));
+ } else {
+ ARROW_RETURN_NOT_OK(
+ Status::Invalid("Unsupported data type: ", src_field->name()));
+ }
+ }
+ }
+ }
+
+ // Build the result
+ assert(sizeof(size_t) >= sizeof(int64_t)); // Make takes signed int64_t for num_rows
+
+ // TODO: check n_rows for cast
+ std::shared_ptr<arrow::RecordBatch> r =
+ arrow::RecordBatch::Make(output_schema, (int64_t)n_rows, arrays);
+ return r;
+ }
+
+ // Returns true if there are no rows
+ bool empty() const { return rows_.empty(); }
+
+ private:
+ // Contains shared_ptr refs for all RecordBatches referred to by the contents of rows_
+ std::unordered_map<uintptr_t, std::shared_ptr<RecordBatch>> _ptr2ref;
+
+ // Row table references
+ std::vector<CompositeReferenceRow<MAX_TABLES>> rows_;
+
+ // Total number of tables in the composite table
+ size_t n_tables_;
+
+ // Adds a RecordBatch ref to the mapping, if needed
+ void add_record_batch_ref(const std::shared_ptr<RecordBatch>& ref) {
+ if (!_ptr2ref.count((uintptr_t)ref.get())) _ptr2ref[(uintptr_t)ref.get()] = ref;
+ }
+
+ template <class Builder, class PrimitiveType>
+ Result<std::shared_ptr<Array>> materialize_primitive_column(size_t i_table,
+ col_index_t i_col) {
+ Builder builder;
+ ARROW_RETURN_NOT_OK(builder.Reserve(rows_.size()));
+ for (row_index_t i_row = 0; i_row < rows_.size(); ++i_row) {
+ const auto& ref = rows_[i_row].refs[i_table];
+ if (ref.batch) {
+ builder.UnsafeAppend(
+ ref.batch->column_data(i_col)->template GetValues<PrimitiveType>(1)[ref.row]);
+ } else {
+ builder.UnsafeAppendNull();
+ }
+ }
+ std::shared_ptr<Array> result;
+ ARROW_RETURN_NOT_OK(builder.Finish(&result));
+ return result;
+ }
+};
+
+class AsofJoinNode : public ExecNode {
+ // Constructs labels for inputs
+ static std::vector<std::string> build_input_labels(
+ const std::vector<ExecNode*>& inputs) {
+ std::vector<std::string> r(inputs.size());
+ for (size_t i = 0; i < r.size(); ++i) r[i] = "input_" + std::to_string(i) + "_label";
+ return r;
+ }
+
+ // Advances the RHS as far as possible to be up to date for the current LHS timestamp
+ bool update_rhs() {
+ auto& lhs = *_state.at(0);
+ auto lhs_latest_time = lhs.get_latest_time();
+ bool any_updated = false;
+ for (size_t i = 1; i < _state.size(); ++i)
+ any_updated |= _state[i]->advance_and_memoize(lhs_latest_time);
+ return any_updated;
+ }
+
+ // Returns false if RHS not up to date for LHS
+ bool is_up_to_date_for_lhs_row() const {
+ auto& lhs = *_state[0];
+ if (lhs.empty()) return false; // can't proceed if nothing on the LHS
+ int64_t lhs_ts = lhs.get_latest_time();
+ for (size_t i = 1; i < _state.size(); ++i) {
+ auto& rhs = *_state[i];
+ if (!rhs.finished()) {
+ // If RHS is finished, then we know it's up to date (but if it isn't, it might be
+ // up to date)
+ if (rhs.empty())
+ return false; // RHS isn't finished, but is empty --> not up to date
+ if (lhs_ts >= rhs.get_latest_time())
+ return false; // TS not up to date (and not finished)
+ }
+ }
+ return true;
+ }
+
+ Result<std::shared_ptr<RecordBatch>> process_inner() {
+ assert(!_state.empty());
+ auto& lhs = *_state.at(0);
+
+ // Construct new target table if needed
+ CompositeReferenceTable<MAX_JOIN_TABLES> dst(_state.size());
+
+ // 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;
+
+ // Advance each of the RHS as far as possible to be up to date for the LHS timestamp
+ bool any_advanced = update_rhs();
+
+ // Only update if we have up-to-date information for the LHS row
+ if (is_up_to_date_for_lhs_row()) {
+ dst.emplace(_state, _options.tolerance);
+ if (!lhs.advance()) break; // if we can't advance LHS, we're done for this batch
+ } else {
+ if ((!any_advanced) && (_state.size() > 1)) break; // need to wait for new data
+ }
+ }
+
+ // Prune memo entries that have expired (to bound memory consumption)
+ if (!lhs.empty()) {
+ for (size_t i = 1; i < _state.size(); ++i) {
+ _state[i]->remove_memo_entries_with_lesser_time(lhs.get_latest_time() -
+ _options.tolerance);
+ }
+ }
+
+ // Emit the batch
+ if (dst.empty()) {
+ return NULLPTR;
+ } else {
+ return dst.materialize(output_schema(), _state);
+ }
+ }
+
+ void process() {
+ std::cerr << "process() begin\n";
+
+ std::lock_guard<std::mutex> guard(_gate);
+ if (finished_.is_finished()) {
+ std::cerr << "InputReceived EARLYEND\n";
+ return;
+ }
+
+ // Process batches while we have data
+ for (;;) {
+ Result<std::shared_ptr<RecordBatch>> result = process_inner();
+
+ if (result.ok()) {
+ auto out_rb = *result;
+ if (!out_rb) break;
+ ++_progress_batches_produced;
+ ExecBatch out_b(*out_rb);
+ outputs_[0]->InputReceived(this, std::move(out_b));
+ } else {
+ StopProducing();
+ ErrorIfNotOk(result.status());
+ return;
+ }
+ }
+
+ std::cerr << "process() end\n";
+
+ // 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 (_state.at(0)->finished()) {
+ total_batches_produced_ = util::make_optional<int>(_progress_batches_produced);
+ StopProducing();
+ assert(total_batches_produced_.has_value());
+ outputs_[0]->InputFinished(this, *total_batches_produced_);
Review Comment:
Can this lead to calling `outputs_[0]->InputFinished` twice? That should only be called once.
##########
cpp/src/arrow/compute/exec/asof_join_node.cc:
##########
@@ -0,0 +1,800 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+#include <iostream>
+#include <unordered_map>
+
+#include <arrow/api.h>
+#include <arrow/compute/api.h>
+#include <arrow/util/optional.h>
+#include "arrow/compute/exec/asof_join.h"
+#include "arrow/compute/exec/exec_plan.h"
+#include "arrow/compute/exec/options.h"
+#include "arrow/compute/exec/schema_util.h"
+#include "arrow/compute/exec/util.h"
+#include "arrow/result.h"
+#include "arrow/status.h"
+#include "arrow/util/checked_cast.h"
+#include "arrow/util/future.h"
+#include "arrow/util/make_unique.h"
+
+#include <condition_variable>
+#include <mutex>
+#include <thread>
+
+namespace arrow {
+namespace compute {
+
+/**
+ * Simple implementation for an unbound concurrent queue
+ */
+template <class T>
+class ConcurrentQueue {
+ public:
+ T pop() {
+ std::unique_lock<std::mutex> lock(mutex_);
+ cond_.wait(lock, [&] { return !queue_.empty(); });
+ auto item = queue_.front();
+ queue_.pop();
+ return item;
+ }
+
+ void push(const T& item) {
+ std::unique_lock<std::mutex> lock(mutex_);
+ queue_.push(item);
+ cond_.notify_one();
+ }
+
+ util::optional<T> try_pop() {
+ // Try to pop the oldest value from the queue (or return nullopt if none)
+ std::unique_lock<std::mutex> lock(mutex_);
+ if (queue_.empty()) {
+ return util::nullopt;
+ } else {
+ auto item = queue_.front();
+ queue_.pop();
+ return item;
+ }
+ }
+
+ bool empty() const {
+ std::unique_lock<std::mutex> lock(mutex_);
+ return queue_.empty();
+ }
+
+ // Un-synchronized access to front
+ // For this to be "safe":
+ // 1) the caller logically guarantees that queue is not empty
+ // 2) pop/try_pop cannot be called concurrently with this
+ const T& unsync_front() const { return queue_.front(); }
+
+ private:
+ std::queue<T> queue_;
+ mutable std::mutex mutex_;
+ std::condition_variable cond_;
+};
+
+struct MemoStore {
+ // Stores last known values for all the keys
+
+ struct Entry {
+ // Timestamp associated with the entry
+ int64_t _time;
+
+ // Batch associated with the entry (perf is probably OK for this; batches change
+ // rarely)
+ std::shared_ptr<arrow::RecordBatch> _batch;
+
+ // Row associated with the entry
+ row_index_t _row;
+ };
+
+ std::unordered_map<KeyType, Entry> _entries;
+
+ void store(const std::shared_ptr<RecordBatch>& batch, row_index_t row, int64_t time,
+ KeyType key) {
+ auto& e = _entries[key];
+ // that we can do this assignment optionally, is why we
+ // can get array 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;
+ }
+
+ util::optional<const Entry*> get_entry_for_key(KeyType key) const {
+ auto e = _entries.find(key);
+ if (_entries.end() == e) return util::nullopt;
+ return util::optional<const Entry*>(&e->second);
+ }
+
+ void remove_entries_with_lesser_time(int64_t ts) {
+ size_t dbg_size0 = _entries.size();
+ for (auto e = _entries.begin(); e != _entries.end();)
+ if (e->second._time < ts)
+ e = _entries.erase(e);
+ else
+ ++e;
+ size_t dbg_size1 = _entries.size();
+ if (dbg_size1 < dbg_size0) {
+ // cerr << "Removed " << dbg_size0-dbg_size1 << " memo entries.\n";
+ }
+ }
+};
+
+class InputState {
+ // InputState correponds to an input
+ // Input record batches are queued up in InputState until processed and
+ // turned into output record batches.
+
+ public:
+ InputState(const std::shared_ptr<arrow::Schema>& schema,
+ const std::string& time_col_name, const std::string& key_col_name,
+ util::optional<KeyType> wildcard_key)
+ : queue_(),
+ wildcard_key_(wildcard_key),
+ schema_(schema),
+ time_col_index_(
+ schema->GetFieldIndex(time_col_name)), // TODO: handle missing field name
+ key_col_index_(schema->GetFieldIndex(key_col_name)) {}
+
+ col_index_t init_src_to_dst_mapping(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 && is_time_or_key_column(i)))
+ src_to_dst_[i] = dst_offset++;
+ return dst_offset;
+ }
+
+ const util::optional<col_index_t>& map_src_to_dst(col_index_t src) const {
+ return src_to_dst_[src];
+ }
+
+ bool is_time_or_key_column(col_index_t i) const {
+ assert(i < schema_->num_fields());
+ return (i == time_col_index_) || (i == key_col_index_);
+ }
+
+ // Gets the latest row index, assuming the queue isn't empty
+ row_index_t get_latest_row() const { return latest_ref_row_; }
+
+ bool empty() const {
+ if (latest_ref_row_ > 0)
+ return false; // cannot be empty if ref row is >0 -- can avoid slow queue lock
+ // below
+ return queue_.empty();
+ }
+
+ int countbatches_processed_() const { return batches_processed_; }
+ int count_total_batches() const { return total_batches_; }
+
+ // Gets latest batch (precondition: must not be empty)
+ const std::shared_ptr<arrow::RecordBatch>& get_latest_batch() const {
+ return queue_.unsync_front();
+ }
+ KeyType get_latest_key() const {
+ return queue_.unsync_front()
+ ->column_data(key_col_index_)
+ ->GetValues<KeyType>(1)[latest_ref_row_];
+ }
+ int64_t get_latest_time() const {
+ return queue_.unsync_front()
+ ->column_data(time_col_index_)
+ ->GetValues<int64_t>(1)[latest_ref_row_];
+ }
+
+ bool finished() const { return batches_processed_ == total_batches_; }
+
+ bool advance() {
+ // 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();
+ if (have_active_batch) {
+ // If we have an active batch
+ if (++latest_ref_row_ >= (row_index_t)queue_.unsync_front()->num_rows()) {
+ // hit the end of the batch, need to get the next batch if possible.
+ ++batches_processed_;
+ latest_ref_row_ = 0;
+ have_active_batch &= !queue_.try_pop();
+ if (have_active_batch)
+ assert(queue_.unsync_front()->num_rows() > 0); // empty batches disallowed
+ }
+ }
+ return have_active_batch;
+ }
+
+ // Advance the data to be immediately past the specified TS, updating latest and
+ // latest_ref_row to the latest data prior to that immediate just past Returns true if
+ // updates were made, false if not.
+ bool advance_and_memoize(int64_t ts) {
+ // Advance the right side row index until we reach the latest right row (for each key)
+ // for the given left timestamp.
+
+ // Check if already updated for TS (or if there is no latest)
+ if (empty()) return false; // can't advance if empty
+ auto latest_time = get_latest_time();
+ if (latest_time > ts) return false; // already advanced
+
+ // Not updated. Try to update and possibly advance.
+ bool updated = false;
+ do {
+ latest_time = get_latest_time();
+ // 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 <= ts) {
+ memo_.store(get_latest_batch(), latest_ref_row_, latest_time, get_latest_key());
+ } else {
+ break; // hit a future timestamp -- done updating for now
+ }
+ updated = true;
+ } while (advance());
+ return updated;
+ }
+
+ void push(const std::shared_ptr<arrow::RecordBatch>& rb) {
+ if (rb->num_rows() > 0) {
+ queue_.push(rb);
+ } else {
+ ++batches_processed_; // don't enqueue empty batches, just record as processed
+ }
+ }
+
+ util::optional<const MemoStore::Entry*> get_memo_entry_for_key(KeyType key) {
+ auto r = memo_.get_entry_for_key(key);
+ if (r.has_value()) return r;
+ if (wildcard_key_.has_value()) r = memo_.get_entry_for_key(*wildcard_key_);
+ return r;
+ }
+
+ util::optional<int64_t> get_memo_time_for_key(KeyType key) {
+ auto r = get_memo_entry_for_key(key);
+ return r.has_value() ? util::make_optional((*r)->_time) : util::nullopt;
+ }
+
+ void remove_memo_entries_with_lesser_time(int64_t ts) {
+ memo_.remove_entries_with_lesser_time(ts);
+ }
+
+ const std::shared_ptr<Schema>& get_schema() const { return schema_; }
+
+ void set_total_batches(int n) {
+ assert(n >= 0);
+ assert(total_batches_ == -1); // shouldn't be set more than once
+ total_batches_ = n;
+ }
+
+ private:
+ // Pending record batches. The latest is the front. Batches cannot be empty.
+ ConcurrentQueue<std::shared_ptr<RecordBatch>> queue_;
+
+ // Wildcard key for this input, if applicable.
+ util::optional<KeyType> wildcard_key_;
+
+ // Schema associated with the input
+ std::shared_ptr<Schema> schema_;
+
+ // Total number of batches (only int because InputFinished uses int)
+ int total_batches_ = -1;
+
+ // Number of batches processed so far (only int because InputFinished uses int)
+ int batches_processed_ = 0;
+
+ // Index of the time col
+ col_index_t time_col_index_;
+
+ // Index of the key col
+ col_index_t key_col_index_;
+
+ // Index of the latest row reference within; if >0 then queue_ cannot be empty
+ row_index_t latest_ref_row_ =
+ 0; // must be < queue_.front()->num_rows() if queue_ is non-empty
+
+ // Stores latest known values for the various keys
+ MemoStore memo_;
+
+ // Mapping of source columns to destination columns
+ std::vector<util::optional<col_index_t>> src_to_dst_;
+};
+
+template <size_t MAX_TABLES>
+struct CompositeReferenceRow {
+ struct Entry {
+ arrow::RecordBatch* batch; // can be NULL if there's no value
+ row_index_t row;
+ };
+ Entry refs[MAX_TABLES];
+};
+
+// A table of composite reference rows. Rows maintain pointers to the
+// constituent record batches, but the overall table retains shared_ptr
+// references to ensure memory remains resident while the table is live.
+//
+// The main reason for this is that, especially for wide tables, joins
+// are effectively row-oriented, rather than column-oriented. Separating
+// the join part from the columnar materialization part simplifies the
+// logic around data types and increases efficiency.
+//
+// We don't put the shared_ptr's into the rows for efficiency reasons.
+template <size_t MAX_TABLES>
+class CompositeReferenceTable {
+ public:
+ explicit CompositeReferenceTable(size_t n_tables) : n_tables_(n_tables) {
+ assert(n_tables_ >= 1);
+ assert(n_tables_ <= MAX_TABLES);
+ }
+
+ size_t n_rows() const { return rows_.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<std::unique_ptr<InputState>>& in, int64_t tolerance) {
+ assert(in.size() == n_tables_);
+
+ // Get the LHS key
+ KeyType key = in[0]->get_latest_key();
+
+ // Add row and setup LHS
+ // (the LHS state comes just from the latest row of the LHS table)
+ assert(!in[0]->empty());
+ const std::shared_ptr<arrow::RecordBatch>& lhs_latest_batch =
+ in[0]->get_latest_batch();
+ row_index_t lhs_latest_row = in[0]->get_latest_row();
+ int64_t lhs_latest_time = in[0]->get_latest_time();
+ 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 = rows_.size() + new_batch_size;
+ if (rows_.capacity() < new_capacity) rows_.reserve(new_capacity);
+ }
+ rows_.resize(rows_.size() + 1);
+ auto& row = rows_.back();
+ row.refs[0].batch = lhs_latest_batch.get();
+ row.refs[0].row = lhs_latest_row;
+ add_record_batch_ref(lhs_latest_batch);
+
+ // 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) {
+ util::optional<const MemoStore::Entry*> opt_entry =
+ in[i]->get_memo_entry_for_key(key);
+ if (opt_entry.has_value()) {
+ assert(*opt_entry);
+ if ((*opt_entry)->_time + tolerance >= lhs_latest_time) {
+ // Have a valid entry
+ const MemoStore::Entry* entry = *opt_entry;
+ row.refs[i].batch = entry->_batch.get();
+ row.refs[i].row = entry->_row;
+ add_record_batch_ref(entry->_batch);
+ continue;
+ }
+ }
+ row.refs[i].batch = NULL;
+ row.refs[i].row = 0;
+ }
+ }
+
+ // Materializes the current reference table into a target record batch
+ Result<std::shared_ptr<RecordBatch>> materialize(
+ const std::shared_ptr<arrow::Schema>& output_schema,
+ const std::vector<std::unique_ptr<InputState>>& state) {
+ // cerr << "materialize BEGIN\n";
+ assert(state.size() == n_tables_);
+ assert(state.size() >= 1);
+
+ // Don't build empty batches
+ size_t n_rows = rows_.size();
+ if (!n_rows) return NULLPTR;
+
+ // Build the arrays column-by-column from the rows
+ std::vector<std::shared_ptr<arrow::Array>> arrays(output_schema->num_fields());
+ for (size_t i_table = 0; i_table < n_tables_; ++i_table) {
+ int n_src_cols = state.at(i_table)->get_schema()->num_fields();
+ {
+ for (col_index_t i_src_col = 0; i_src_col < n_src_cols; ++i_src_col) {
+ util::optional<col_index_t> i_dst_col_opt =
+ state[i_table]->map_src_to_dst(i_src_col);
+ if (!i_dst_col_opt) continue;
+ col_index_t i_dst_col = *i_dst_col_opt;
+ const auto& src_field = state[i_table]->get_schema()->field(i_src_col);
+ const auto& dst_field = output_schema->field(i_dst_col);
+ assert(src_field->type()->Equals(dst_field->type()));
+ assert(src_field->name() == dst_field->name());
+ const auto& field_type = src_field->type();
+
+ if (field_type->Equals(arrow::int32())) {
+ ARROW_ASSIGN_OR_RAISE(
+ arrays.at(i_dst_col),
+ (materialize_primitive_column<arrow::Int32Builder, int32_t>(i_table,
+ i_src_col)));
+ } else if (field_type->Equals(arrow::int64())) {
+ ARROW_ASSIGN_OR_RAISE(
+ arrays.at(i_dst_col),
+ (materialize_primitive_column<arrow::Int64Builder, int64_t>(i_table,
+ i_src_col)));
+ } else if (field_type->Equals(arrow::float64())) {
+ ARROW_ASSIGN_OR_RAISE(
+ arrays.at(i_dst_col),
+ (materialize_primitive_column<arrow::DoubleBuilder, double>(i_table,
+ i_src_col)));
+ } else {
+ ARROW_RETURN_NOT_OK(
+ Status::Invalid("Unsupported data type: ", src_field->name()));
+ }
+ }
+ }
+ }
+
+ // Build the result
+ assert(sizeof(size_t) >= sizeof(int64_t)); // Make takes signed int64_t for num_rows
+
+ // TODO: check n_rows for cast
+ std::shared_ptr<arrow::RecordBatch> r =
+ arrow::RecordBatch::Make(output_schema, (int64_t)n_rows, arrays);
+ return r;
+ }
+
+ // Returns true if there are no rows
+ bool empty() const { return rows_.empty(); }
+
+ private:
+ // Contains shared_ptr refs for all RecordBatches referred to by the contents of rows_
+ std::unordered_map<uintptr_t, std::shared_ptr<RecordBatch>> _ptr2ref;
+
+ // Row table references
+ std::vector<CompositeReferenceRow<MAX_TABLES>> rows_;
+
+ // Total number of tables in the composite table
+ size_t n_tables_;
+
+ // Adds a RecordBatch ref to the mapping, if needed
+ void add_record_batch_ref(const std::shared_ptr<RecordBatch>& ref) {
+ if (!_ptr2ref.count((uintptr_t)ref.get())) _ptr2ref[(uintptr_t)ref.get()] = ref;
+ }
+
+ template <class Builder, class PrimitiveType>
+ Result<std::shared_ptr<Array>> materialize_primitive_column(size_t i_table,
+ col_index_t i_col) {
+ Builder builder;
+ ARROW_RETURN_NOT_OK(builder.Reserve(rows_.size()));
+ for (row_index_t i_row = 0; i_row < rows_.size(); ++i_row) {
+ const auto& ref = rows_[i_row].refs[i_table];
+ if (ref.batch) {
+ builder.UnsafeAppend(
+ ref.batch->column_data(i_col)->template GetValues<PrimitiveType>(1)[ref.row]);
+ } else {
+ builder.UnsafeAppendNull();
+ }
+ }
+ std::shared_ptr<Array> result;
+ ARROW_RETURN_NOT_OK(builder.Finish(&result));
+ return result;
+ }
+};
+
+class AsofJoinNode : public ExecNode {
+ // Constructs labels for inputs
+ static std::vector<std::string> build_input_labels(
+ const std::vector<ExecNode*>& inputs) {
+ std::vector<std::string> r(inputs.size());
+ for (size_t i = 0; i < r.size(); ++i) r[i] = "input_" + std::to_string(i) + "_label";
+ return r;
+ }
+
+ // Advances the RHS as far as possible to be up to date for the current LHS timestamp
+ bool update_rhs() {
+ auto& lhs = *_state.at(0);
+ auto lhs_latest_time = lhs.get_latest_time();
+ bool any_updated = false;
+ for (size_t i = 1; i < _state.size(); ++i)
+ any_updated |= _state[i]->advance_and_memoize(lhs_latest_time);
+ return any_updated;
+ }
+
+ // Returns false if RHS not up to date for LHS
+ bool is_up_to_date_for_lhs_row() const {
+ auto& lhs = *_state[0];
+ if (lhs.empty()) return false; // can't proceed if nothing on the LHS
+ int64_t lhs_ts = lhs.get_latest_time();
+ for (size_t i = 1; i < _state.size(); ++i) {
+ auto& rhs = *_state[i];
+ if (!rhs.finished()) {
+ // If RHS is finished, then we know it's up to date (but if it isn't, it might be
+ // up to date)
+ if (rhs.empty())
+ return false; // RHS isn't finished, but is empty --> not up to date
+ if (lhs_ts >= rhs.get_latest_time())
+ return false; // TS not up to date (and not finished)
+ }
+ }
+ return true;
+ }
+
+ Result<std::shared_ptr<RecordBatch>> process_inner() {
+ assert(!_state.empty());
+ auto& lhs = *_state.at(0);
+
+ // Construct new target table if needed
+ CompositeReferenceTable<MAX_JOIN_TABLES> dst(_state.size());
+
+ // 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;
+
+ // Advance each of the RHS as far as possible to be up to date for the LHS timestamp
+ bool any_advanced = update_rhs();
+
+ // Only update if we have up-to-date information for the LHS row
+ if (is_up_to_date_for_lhs_row()) {
+ dst.emplace(_state, _options.tolerance);
+ if (!lhs.advance()) break; // if we can't advance LHS, we're done for this batch
+ } else {
+ if ((!any_advanced) && (_state.size() > 1)) break; // need to wait for new data
+ }
+ }
+
+ // Prune memo entries that have expired (to bound memory consumption)
+ if (!lhs.empty()) {
+ for (size_t i = 1; i < _state.size(); ++i) {
+ _state[i]->remove_memo_entries_with_lesser_time(lhs.get_latest_time() -
+ _options.tolerance);
+ }
+ }
+
+ // Emit the batch
+ if (dst.empty()) {
+ return NULLPTR;
+ } else {
+ return dst.materialize(output_schema(), _state);
+ }
+ }
+
+ void process() {
+ std::cerr << "process() begin\n";
+
+ std::lock_guard<std::mutex> guard(_gate);
+ if (finished_.is_finished()) {
+ std::cerr << "InputReceived EARLYEND\n";
+ return;
+ }
+
+ // Process batches while we have data
+ for (;;) {
+ Result<std::shared_ptr<RecordBatch>> result = process_inner();
+
+ if (result.ok()) {
+ auto out_rb = *result;
+ if (!out_rb) break;
+ ++_progress_batches_produced;
+ ExecBatch out_b(*out_rb);
+ outputs_[0]->InputReceived(this, std::move(out_b));
+ } else {
+ StopProducing();
+ ErrorIfNotOk(result.status());
+ return;
+ }
+ }
+
+ std::cerr << "process() end\n";
+
+ // 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 (_state.at(0)->finished()) {
+ total_batches_produced_ = util::make_optional<int>(_progress_batches_produced);
+ StopProducing();
+ assert(total_batches_produced_.has_value());
+ outputs_[0]->InputFinished(this, *total_batches_produced_);
+ }
+ }
+
+ void process_thread() {
+ std::cerr << "AsofJoinNode::process_thread started.\n";
+ for (;;) {
+ if (!_process.pop()) {
+ std::cerr << "AsofJoinNode::process_thread done.\n";
+ return;
+ }
+ process();
+ }
+ }
+
+ static void process_thread_wrapper(AsofJoinNode* node) { node->process_thread(); }
+
+ public:
+ AsofJoinNode(ExecPlan* plan, NodeVector inputs, std::vector<std::string> input_labels,
+ const AsofJoinNodeOptions& join_options,
+ std::shared_ptr<Schema> output_schema,
+ std::unique_ptr<AsofJoinSchema> schema_mgr);
+
+ virtual ~AsofJoinNode() {
+ _process.push(false); // poison pill
+ _process_thread.join();
+ }
+
+ static arrow::Result<ExecNode*> Make(ExecPlan* plan, std::vector<ExecNode*> inputs,
+ const ExecNodeOptions& options) {
+ std::unique_ptr<AsofJoinSchema> schema_mgr =
+ ::arrow::internal::make_unique<AsofJoinSchema>();
+
+ const auto& join_options = checked_cast<const AsofJoinNodeOptions&>(options);
+ std::shared_ptr<Schema> output_schema =
+ schema_mgr->MakeOutputSchema(inputs, join_options);
+
+ std::vector<std::string> input_labels(inputs.size());
Review Comment:
Can you validate there are at least two inputs here?
##########
cpp/src/arrow/compute/exec/asof_join_node.cc:
##########
@@ -0,0 +1,800 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+#include <iostream>
+#include <unordered_map>
+
+#include <arrow/api.h>
+#include <arrow/compute/api.h>
+#include <arrow/util/optional.h>
+#include "arrow/compute/exec/asof_join.h"
+#include "arrow/compute/exec/exec_plan.h"
+#include "arrow/compute/exec/options.h"
+#include "arrow/compute/exec/schema_util.h"
+#include "arrow/compute/exec/util.h"
+#include "arrow/result.h"
+#include "arrow/status.h"
+#include "arrow/util/checked_cast.h"
+#include "arrow/util/future.h"
+#include "arrow/util/make_unique.h"
+
+#include <condition_variable>
+#include <mutex>
+#include <thread>
+
+namespace arrow {
+namespace compute {
+
+/**
+ * Simple implementation for an unbound concurrent queue
+ */
+template <class T>
+class ConcurrentQueue {
+ public:
+ T pop() {
+ std::unique_lock<std::mutex> lock(mutex_);
+ cond_.wait(lock, [&] { return !queue_.empty(); });
+ auto item = queue_.front();
+ queue_.pop();
+ return item;
+ }
+
+ void push(const T& item) {
+ std::unique_lock<std::mutex> lock(mutex_);
+ queue_.push(item);
+ cond_.notify_one();
+ }
+
+ util::optional<T> try_pop() {
+ // Try to pop the oldest value from the queue (or return nullopt if none)
+ std::unique_lock<std::mutex> lock(mutex_);
+ if (queue_.empty()) {
+ return util::nullopt;
+ } else {
+ auto item = queue_.front();
+ queue_.pop();
+ return item;
+ }
+ }
+
+ bool empty() const {
+ std::unique_lock<std::mutex> lock(mutex_);
+ return queue_.empty();
+ }
+
+ // Un-synchronized access to front
+ // For this to be "safe":
+ // 1) the caller logically guarantees that queue is not empty
+ // 2) pop/try_pop cannot be called concurrently with this
+ const T& unsync_front() const { return queue_.front(); }
+
+ private:
+ std::queue<T> queue_;
+ mutable std::mutex mutex_;
+ std::condition_variable cond_;
+};
+
+struct MemoStore {
+ // Stores last known values for all the keys
+
+ struct Entry {
+ // Timestamp associated with the entry
+ int64_t _time;
+
+ // Batch associated with the entry (perf is probably OK for this; batches change
+ // rarely)
+ std::shared_ptr<arrow::RecordBatch> _batch;
+
+ // Row associated with the entry
+ row_index_t _row;
+ };
+
+ std::unordered_map<KeyType, Entry> _entries;
+
+ void store(const std::shared_ptr<RecordBatch>& batch, row_index_t row, int64_t time,
+ KeyType key) {
+ auto& e = _entries[key];
+ // that we can do this assignment optionally, is why we
+ // can get array 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;
+ }
+
+ util::optional<const Entry*> get_entry_for_key(KeyType key) const {
+ auto e = _entries.find(key);
+ if (_entries.end() == e) return util::nullopt;
+ return util::optional<const Entry*>(&e->second);
+ }
+
+ void remove_entries_with_lesser_time(int64_t ts) {
+ size_t dbg_size0 = _entries.size();
+ for (auto e = _entries.begin(); e != _entries.end();)
+ if (e->second._time < ts)
+ e = _entries.erase(e);
+ else
+ ++e;
+ size_t dbg_size1 = _entries.size();
+ if (dbg_size1 < dbg_size0) {
+ // cerr << "Removed " << dbg_size0-dbg_size1 << " memo entries.\n";
+ }
+ }
+};
+
+class InputState {
+ // InputState correponds to an input
+ // Input record batches are queued up in InputState until processed and
+ // turned into output record batches.
+
+ public:
+ InputState(const std::shared_ptr<arrow::Schema>& schema,
+ const std::string& time_col_name, const std::string& key_col_name,
+ util::optional<KeyType> wildcard_key)
+ : queue_(),
+ wildcard_key_(wildcard_key),
+ schema_(schema),
+ time_col_index_(
+ schema->GetFieldIndex(time_col_name)), // TODO: handle missing field name
+ key_col_index_(schema->GetFieldIndex(key_col_name)) {}
+
+ col_index_t init_src_to_dst_mapping(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 && is_time_or_key_column(i)))
+ src_to_dst_[i] = dst_offset++;
+ return dst_offset;
+ }
+
+ const util::optional<col_index_t>& map_src_to_dst(col_index_t src) const {
+ return src_to_dst_[src];
+ }
+
+ bool is_time_or_key_column(col_index_t i) const {
+ assert(i < schema_->num_fields());
+ return (i == time_col_index_) || (i == key_col_index_);
+ }
+
+ // Gets the latest row index, assuming the queue isn't empty
+ row_index_t get_latest_row() const { return latest_ref_row_; }
+
+ bool empty() const {
+ if (latest_ref_row_ > 0)
+ return false; // cannot be empty if ref row is >0 -- can avoid slow queue lock
+ // below
+ return queue_.empty();
+ }
+
+ int countbatches_processed_() const { return batches_processed_; }
+ int count_total_batches() const { return total_batches_; }
+
+ // Gets latest batch (precondition: must not be empty)
+ const std::shared_ptr<arrow::RecordBatch>& get_latest_batch() const {
+ return queue_.unsync_front();
+ }
+ KeyType get_latest_key() const {
+ return queue_.unsync_front()
+ ->column_data(key_col_index_)
+ ->GetValues<KeyType>(1)[latest_ref_row_];
+ }
+ int64_t get_latest_time() const {
+ return queue_.unsync_front()
+ ->column_data(time_col_index_)
+ ->GetValues<int64_t>(1)[latest_ref_row_];
+ }
+
+ bool finished() const { return batches_processed_ == total_batches_; }
+
+ bool advance() {
+ // 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();
+ if (have_active_batch) {
+ // If we have an active batch
+ if (++latest_ref_row_ >= (row_index_t)queue_.unsync_front()->num_rows()) {
+ // hit the end of the batch, need to get the next batch if possible.
+ ++batches_processed_;
+ latest_ref_row_ = 0;
+ have_active_batch &= !queue_.try_pop();
+ if (have_active_batch)
+ assert(queue_.unsync_front()->num_rows() > 0); // empty batches disallowed
+ }
+ }
+ return have_active_batch;
+ }
+
+ // Advance the data to be immediately past the specified TS, updating latest and
+ // latest_ref_row to the latest data prior to that immediate just past Returns true if
+ // updates were made, false if not.
+ bool advance_and_memoize(int64_t ts) {
+ // Advance the right side row index until we reach the latest right row (for each key)
+ // for the given left timestamp.
+
+ // Check if already updated for TS (or if there is no latest)
+ if (empty()) return false; // can't advance if empty
+ auto latest_time = get_latest_time();
+ if (latest_time > ts) return false; // already advanced
+
+ // Not updated. Try to update and possibly advance.
+ bool updated = false;
+ do {
+ latest_time = get_latest_time();
+ // 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 <= ts) {
+ memo_.store(get_latest_batch(), latest_ref_row_, latest_time, get_latest_key());
+ } else {
+ break; // hit a future timestamp -- done updating for now
+ }
+ updated = true;
+ } while (advance());
+ return updated;
+ }
+
+ void push(const std::shared_ptr<arrow::RecordBatch>& rb) {
+ if (rb->num_rows() > 0) {
+ queue_.push(rb);
+ } else {
+ ++batches_processed_; // don't enqueue empty batches, just record as processed
+ }
+ }
+
+ util::optional<const MemoStore::Entry*> get_memo_entry_for_key(KeyType key) {
+ auto r = memo_.get_entry_for_key(key);
+ if (r.has_value()) return r;
+ if (wildcard_key_.has_value()) r = memo_.get_entry_for_key(*wildcard_key_);
+ return r;
+ }
+
+ util::optional<int64_t> get_memo_time_for_key(KeyType key) {
+ auto r = get_memo_entry_for_key(key);
+ return r.has_value() ? util::make_optional((*r)->_time) : util::nullopt;
Review Comment:
```suggestion
return r.has_value() ? (*r)->_time : util::nullopt;
```
Minor nit: elsewhere we just rely on the implicit constructor of `util::optional`
##########
cpp/src/arrow/compute/exec/asof_join_node.cc:
##########
@@ -0,0 +1,800 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+#include <iostream>
+#include <unordered_map>
+
+#include <arrow/api.h>
+#include <arrow/compute/api.h>
+#include <arrow/util/optional.h>
+#include "arrow/compute/exec/asof_join.h"
+#include "arrow/compute/exec/exec_plan.h"
+#include "arrow/compute/exec/options.h"
+#include "arrow/compute/exec/schema_util.h"
+#include "arrow/compute/exec/util.h"
+#include "arrow/result.h"
+#include "arrow/status.h"
+#include "arrow/util/checked_cast.h"
+#include "arrow/util/future.h"
+#include "arrow/util/make_unique.h"
+
+#include <condition_variable>
+#include <mutex>
+#include <thread>
+
+namespace arrow {
+namespace compute {
+
+/**
+ * Simple implementation for an unbound concurrent queue
+ */
+template <class T>
+class ConcurrentQueue {
+ public:
+ T pop() {
+ std::unique_lock<std::mutex> lock(mutex_);
+ cond_.wait(lock, [&] { return !queue_.empty(); });
+ auto item = queue_.front();
+ queue_.pop();
+ return item;
+ }
+
+ void push(const T& item) {
+ std::unique_lock<std::mutex> lock(mutex_);
+ queue_.push(item);
+ cond_.notify_one();
+ }
+
+ util::optional<T> try_pop() {
+ // Try to pop the oldest value from the queue (or return nullopt if none)
+ std::unique_lock<std::mutex> lock(mutex_);
+ if (queue_.empty()) {
+ return util::nullopt;
+ } else {
+ auto item = queue_.front();
+ queue_.pop();
+ return item;
+ }
+ }
+
+ bool empty() const {
+ std::unique_lock<std::mutex> lock(mutex_);
+ return queue_.empty();
+ }
+
+ // Un-synchronized access to front
+ // For this to be "safe":
+ // 1) the caller logically guarantees that queue is not empty
+ // 2) pop/try_pop cannot be called concurrently with this
+ const T& unsync_front() const { return queue_.front(); }
+
+ private:
+ std::queue<T> queue_;
+ mutable std::mutex mutex_;
+ std::condition_variable cond_;
+};
+
+struct MemoStore {
+ // Stores last known values for all the keys
+
+ struct Entry {
+ // Timestamp associated with the entry
+ int64_t _time;
+
+ // Batch associated with the entry (perf is probably OK for this; batches change
+ // rarely)
+ std::shared_ptr<arrow::RecordBatch> _batch;
+
+ // Row associated with the entry
+ row_index_t _row;
+ };
+
+ std::unordered_map<KeyType, Entry> _entries;
+
+ void store(const std::shared_ptr<RecordBatch>& batch, row_index_t row, int64_t time,
+ KeyType key) {
+ auto& e = _entries[key];
+ // that we can do this assignment optionally, is why we
+ // can get array 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;
+ }
+
+ util::optional<const Entry*> get_entry_for_key(KeyType key) const {
+ auto e = _entries.find(key);
+ if (_entries.end() == e) return util::nullopt;
+ return util::optional<const Entry*>(&e->second);
+ }
+
+ void remove_entries_with_lesser_time(int64_t ts) {
+ size_t dbg_size0 = _entries.size();
+ for (auto e = _entries.begin(); e != _entries.end();)
+ if (e->second._time < ts)
+ e = _entries.erase(e);
+ else
+ ++e;
+ size_t dbg_size1 = _entries.size();
+ if (dbg_size1 < dbg_size0) {
+ // cerr << "Removed " << dbg_size0-dbg_size1 << " memo entries.\n";
+ }
+ }
+};
+
+class InputState {
+ // InputState correponds to an input
+ // Input record batches are queued up in InputState until processed and
+ // turned into output record batches.
+
+ public:
+ InputState(const std::shared_ptr<arrow::Schema>& schema,
+ const std::string& time_col_name, const std::string& key_col_name,
+ util::optional<KeyType> wildcard_key)
+ : queue_(),
+ wildcard_key_(wildcard_key),
+ schema_(schema),
+ time_col_index_(
+ schema->GetFieldIndex(time_col_name)), // TODO: handle missing field name
+ key_col_index_(schema->GetFieldIndex(key_col_name)) {}
+
+ col_index_t init_src_to_dst_mapping(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 && is_time_or_key_column(i)))
+ src_to_dst_[i] = dst_offset++;
+ return dst_offset;
+ }
+
+ const util::optional<col_index_t>& map_src_to_dst(col_index_t src) const {
+ return src_to_dst_[src];
+ }
+
+ bool is_time_or_key_column(col_index_t i) const {
+ assert(i < schema_->num_fields());
+ return (i == time_col_index_) || (i == key_col_index_);
+ }
+
+ // Gets the latest row index, assuming the queue isn't empty
+ row_index_t get_latest_row() const { return latest_ref_row_; }
+
+ bool empty() const {
+ if (latest_ref_row_ > 0)
+ return false; // cannot be empty if ref row is >0 -- can avoid slow queue lock
+ // below
+ return queue_.empty();
+ }
+
+ int countbatches_processed_() const { return batches_processed_; }
+ int count_total_batches() const { return total_batches_; }
+
+ // Gets latest batch (precondition: must not be empty)
+ const std::shared_ptr<arrow::RecordBatch>& get_latest_batch() const {
+ return queue_.unsync_front();
+ }
+ KeyType get_latest_key() const {
+ return queue_.unsync_front()
+ ->column_data(key_col_index_)
+ ->GetValues<KeyType>(1)[latest_ref_row_];
+ }
+ int64_t get_latest_time() const {
+ return queue_.unsync_front()
+ ->column_data(time_col_index_)
+ ->GetValues<int64_t>(1)[latest_ref_row_];
+ }
+
+ bool finished() const { return batches_processed_ == total_batches_; }
+
+ bool advance() {
+ // 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();
+ if (have_active_batch) {
+ // If we have an active batch
+ if (++latest_ref_row_ >= (row_index_t)queue_.unsync_front()->num_rows()) {
+ // hit the end of the batch, need to get the next batch if possible.
+ ++batches_processed_;
+ latest_ref_row_ = 0;
+ have_active_batch &= !queue_.try_pop();
+ if (have_active_batch)
+ assert(queue_.unsync_front()->num_rows() > 0); // empty batches disallowed
+ }
+ }
+ return have_active_batch;
+ }
+
+ // Advance the data to be immediately past the specified TS, updating latest and
+ // latest_ref_row to the latest data prior to that immediate just past Returns true if
+ // updates were made, false if not.
Review Comment:
Can you reword this description? I'm not sure I understand `to the latest data prior to that immediate just past`
##########
cpp/src/arrow/compute/exec/asof_join_node.cc:
##########
@@ -0,0 +1,800 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+#include <iostream>
+#include <unordered_map>
+
+#include <arrow/api.h>
+#include <arrow/compute/api.h>
+#include <arrow/util/optional.h>
+#include "arrow/compute/exec/asof_join.h"
+#include "arrow/compute/exec/exec_plan.h"
+#include "arrow/compute/exec/options.h"
+#include "arrow/compute/exec/schema_util.h"
+#include "arrow/compute/exec/util.h"
+#include "arrow/result.h"
+#include "arrow/status.h"
+#include "arrow/util/checked_cast.h"
+#include "arrow/util/future.h"
+#include "arrow/util/make_unique.h"
+
+#include <condition_variable>
+#include <mutex>
+#include <thread>
+
+namespace arrow {
+namespace compute {
+
+/**
+ * Simple implementation for an unbound concurrent queue
+ */
+template <class T>
+class ConcurrentQueue {
+ public:
+ T pop() {
+ std::unique_lock<std::mutex> lock(mutex_);
+ cond_.wait(lock, [&] { return !queue_.empty(); });
+ auto item = queue_.front();
+ queue_.pop();
+ return item;
+ }
+
+ void push(const T& item) {
+ std::unique_lock<std::mutex> lock(mutex_);
+ queue_.push(item);
+ cond_.notify_one();
+ }
+
+ util::optional<T> try_pop() {
+ // Try to pop the oldest value from the queue (or return nullopt if none)
+ std::unique_lock<std::mutex> lock(mutex_);
+ if (queue_.empty()) {
+ return util::nullopt;
+ } else {
+ auto item = queue_.front();
+ queue_.pop();
+ return item;
+ }
+ }
+
+ bool empty() const {
+ std::unique_lock<std::mutex> lock(mutex_);
+ return queue_.empty();
+ }
+
+ // Un-synchronized access to front
+ // For this to be "safe":
+ // 1) the caller logically guarantees that queue is not empty
+ // 2) pop/try_pop cannot be called concurrently with this
+ const T& unsync_front() const { return queue_.front(); }
+
+ private:
+ std::queue<T> queue_;
+ mutable std::mutex mutex_;
+ std::condition_variable cond_;
+};
+
+struct MemoStore {
+ // Stores last known values for all the keys
+
+ struct Entry {
+ // Timestamp associated with the entry
+ int64_t _time;
+
+ // Batch associated with the entry (perf is probably OK for this; batches change
+ // rarely)
+ std::shared_ptr<arrow::RecordBatch> _batch;
+
+ // Row associated with the entry
+ row_index_t _row;
+ };
+
+ std::unordered_map<KeyType, Entry> _entries;
+
+ void store(const std::shared_ptr<RecordBatch>& batch, row_index_t row, int64_t time,
+ KeyType key) {
+ auto& e = _entries[key];
+ // that we can do this assignment optionally, is why we
+ // can get array 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;
+ }
+
+ util::optional<const Entry*> get_entry_for_key(KeyType key) const {
+ auto e = _entries.find(key);
+ if (_entries.end() == e) return util::nullopt;
+ return util::optional<const Entry*>(&e->second);
+ }
+
+ void remove_entries_with_lesser_time(int64_t ts) {
+ size_t dbg_size0 = _entries.size();
+ for (auto e = _entries.begin(); e != _entries.end();)
+ if (e->second._time < ts)
+ e = _entries.erase(e);
+ else
+ ++e;
+ size_t dbg_size1 = _entries.size();
+ if (dbg_size1 < dbg_size0) {
+ // cerr << "Removed " << dbg_size0-dbg_size1 << " memo entries.\n";
+ }
+ }
+};
+
+class InputState {
+ // InputState correponds to an input
+ // Input record batches are queued up in InputState until processed and
+ // turned into output record batches.
+
+ public:
+ InputState(const std::shared_ptr<arrow::Schema>& schema,
+ const std::string& time_col_name, const std::string& key_col_name,
+ util::optional<KeyType> wildcard_key)
+ : queue_(),
+ wildcard_key_(wildcard_key),
+ schema_(schema),
+ time_col_index_(
+ schema->GetFieldIndex(time_col_name)), // TODO: handle missing field name
+ key_col_index_(schema->GetFieldIndex(key_col_name)) {}
+
+ col_index_t init_src_to_dst_mapping(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 && is_time_or_key_column(i)))
+ src_to_dst_[i] = dst_offset++;
+ return dst_offset;
+ }
+
+ const util::optional<col_index_t>& map_src_to_dst(col_index_t src) const {
+ return src_to_dst_[src];
+ }
+
+ bool is_time_or_key_column(col_index_t i) const {
+ assert(i < schema_->num_fields());
+ return (i == time_col_index_) || (i == key_col_index_);
+ }
+
+ // Gets the latest row index, assuming the queue isn't empty
+ row_index_t get_latest_row() const { return latest_ref_row_; }
+
+ bool empty() const {
+ if (latest_ref_row_ > 0)
+ return false; // cannot be empty if ref row is >0 -- can avoid slow queue lock
+ // below
+ return queue_.empty();
+ }
+
+ int countbatches_processed_() const { return batches_processed_; }
+ int count_total_batches() const { return total_batches_; }
+
+ // Gets latest batch (precondition: must not be empty)
+ const std::shared_ptr<arrow::RecordBatch>& get_latest_batch() const {
+ return queue_.unsync_front();
+ }
+ KeyType get_latest_key() const {
+ return queue_.unsync_front()
+ ->column_data(key_col_index_)
+ ->GetValues<KeyType>(1)[latest_ref_row_];
+ }
+ int64_t get_latest_time() const {
+ return queue_.unsync_front()
+ ->column_data(time_col_index_)
+ ->GetValues<int64_t>(1)[latest_ref_row_];
+ }
+
+ bool finished() const { return batches_processed_ == total_batches_; }
+
+ bool advance() {
+ // 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();
+ if (have_active_batch) {
Review Comment:
```suggestion
if (!empty()) {
```
##########
cpp/src/arrow/compute/exec/asof_join_node.cc:
##########
@@ -0,0 +1,800 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+#include <iostream>
+#include <unordered_map>
+
+#include <arrow/api.h>
+#include <arrow/compute/api.h>
+#include <arrow/util/optional.h>
+#include "arrow/compute/exec/asof_join.h"
+#include "arrow/compute/exec/exec_plan.h"
+#include "arrow/compute/exec/options.h"
+#include "arrow/compute/exec/schema_util.h"
+#include "arrow/compute/exec/util.h"
+#include "arrow/result.h"
+#include "arrow/status.h"
+#include "arrow/util/checked_cast.h"
+#include "arrow/util/future.h"
+#include "arrow/util/make_unique.h"
+
+#include <condition_variable>
+#include <mutex>
+#include <thread>
+
+namespace arrow {
+namespace compute {
+
+/**
+ * Simple implementation for an unbound concurrent queue
+ */
+template <class T>
+class ConcurrentQueue {
+ public:
+ T pop() {
+ std::unique_lock<std::mutex> lock(mutex_);
+ cond_.wait(lock, [&] { return !queue_.empty(); });
+ auto item = queue_.front();
+ queue_.pop();
+ return item;
+ }
+
+ void push(const T& item) {
+ std::unique_lock<std::mutex> lock(mutex_);
+ queue_.push(item);
+ cond_.notify_one();
+ }
+
+ util::optional<T> try_pop() {
+ // Try to pop the oldest value from the queue (or return nullopt if none)
+ std::unique_lock<std::mutex> lock(mutex_);
+ if (queue_.empty()) {
+ return util::nullopt;
+ } else {
+ auto item = queue_.front();
+ queue_.pop();
+ return item;
+ }
+ }
+
+ bool empty() const {
+ std::unique_lock<std::mutex> lock(mutex_);
+ return queue_.empty();
+ }
+
+ // Un-synchronized access to front
+ // For this to be "safe":
+ // 1) the caller logically guarantees that queue is not empty
+ // 2) pop/try_pop cannot be called concurrently with this
+ const T& unsync_front() const { return queue_.front(); }
+
+ private:
+ std::queue<T> queue_;
+ mutable std::mutex mutex_;
+ std::condition_variable cond_;
+};
+
+struct MemoStore {
+ // Stores last known values for all the keys
+
+ struct Entry {
+ // Timestamp associated with the entry
+ int64_t _time;
+
+ // Batch associated with the entry (perf is probably OK for this; batches change
+ // rarely)
+ std::shared_ptr<arrow::RecordBatch> _batch;
+
+ // Row associated with the entry
+ row_index_t _row;
+ };
+
+ std::unordered_map<KeyType, Entry> _entries;
+
+ void store(const std::shared_ptr<RecordBatch>& batch, row_index_t row, int64_t time,
+ KeyType key) {
+ auto& e = _entries[key];
+ // that we can do this assignment optionally, is why we
+ // can get array 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;
+ }
+
+ util::optional<const Entry*> get_entry_for_key(KeyType key) const {
+ auto e = _entries.find(key);
+ if (_entries.end() == e) return util::nullopt;
+ return util::optional<const Entry*>(&e->second);
+ }
+
+ void remove_entries_with_lesser_time(int64_t ts) {
+ size_t dbg_size0 = _entries.size();
+ for (auto e = _entries.begin(); e != _entries.end();)
+ if (e->second._time < ts)
+ e = _entries.erase(e);
+ else
+ ++e;
+ size_t dbg_size1 = _entries.size();
+ if (dbg_size1 < dbg_size0) {
+ // cerr << "Removed " << dbg_size0-dbg_size1 << " memo entries.\n";
+ }
+ }
+};
+
+class InputState {
+ // InputState correponds to an input
+ // Input record batches are queued up in InputState until processed and
+ // turned into output record batches.
+
+ public:
+ InputState(const std::shared_ptr<arrow::Schema>& schema,
+ const std::string& time_col_name, const std::string& key_col_name,
+ util::optional<KeyType> wildcard_key)
+ : queue_(),
+ wildcard_key_(wildcard_key),
+ schema_(schema),
+ time_col_index_(
+ schema->GetFieldIndex(time_col_name)), // TODO: handle missing field name
+ key_col_index_(schema->GetFieldIndex(key_col_name)) {}
+
+ col_index_t init_src_to_dst_mapping(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 && is_time_or_key_column(i)))
+ src_to_dst_[i] = dst_offset++;
+ return dst_offset;
+ }
+
+ const util::optional<col_index_t>& map_src_to_dst(col_index_t src) const {
+ return src_to_dst_[src];
+ }
+
+ bool is_time_or_key_column(col_index_t i) const {
+ assert(i < schema_->num_fields());
Review Comment:
```suggestion
DCHECK_LT(i, schema_->num_fields());
```
##########
cpp/src/arrow/compute/exec/asof_join_node.cc:
##########
@@ -0,0 +1,800 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+#include <iostream>
+#include <unordered_map>
+
+#include <arrow/api.h>
+#include <arrow/compute/api.h>
+#include <arrow/util/optional.h>
+#include "arrow/compute/exec/asof_join.h"
+#include "arrow/compute/exec/exec_plan.h"
+#include "arrow/compute/exec/options.h"
+#include "arrow/compute/exec/schema_util.h"
+#include "arrow/compute/exec/util.h"
+#include "arrow/result.h"
+#include "arrow/status.h"
+#include "arrow/util/checked_cast.h"
+#include "arrow/util/future.h"
+#include "arrow/util/make_unique.h"
+
+#include <condition_variable>
+#include <mutex>
+#include <thread>
+
+namespace arrow {
+namespace compute {
+
+/**
+ * Simple implementation for an unbound concurrent queue
+ */
+template <class T>
+class ConcurrentQueue {
+ public:
+ T pop() {
+ std::unique_lock<std::mutex> lock(mutex_);
+ cond_.wait(lock, [&] { return !queue_.empty(); });
+ auto item = queue_.front();
+ queue_.pop();
+ return item;
+ }
+
+ void push(const T& item) {
+ std::unique_lock<std::mutex> lock(mutex_);
+ queue_.push(item);
+ cond_.notify_one();
+ }
+
+ util::optional<T> try_pop() {
+ // Try to pop the oldest value from the queue (or return nullopt if none)
+ std::unique_lock<std::mutex> lock(mutex_);
+ if (queue_.empty()) {
+ return util::nullopt;
+ } else {
+ auto item = queue_.front();
+ queue_.pop();
+ return item;
+ }
+ }
+
+ bool empty() const {
+ std::unique_lock<std::mutex> lock(mutex_);
+ return queue_.empty();
+ }
+
+ // Un-synchronized access to front
+ // For this to be "safe":
+ // 1) the caller logically guarantees that queue is not empty
+ // 2) pop/try_pop cannot be called concurrently with this
+ const T& unsync_front() const { return queue_.front(); }
+
+ private:
+ std::queue<T> queue_;
+ mutable std::mutex mutex_;
+ std::condition_variable cond_;
+};
+
+struct MemoStore {
+ // Stores last known values for all the keys
+
+ struct Entry {
+ // Timestamp associated with the entry
+ int64_t _time;
+
+ // Batch associated with the entry (perf is probably OK for this; batches change
+ // rarely)
+ std::shared_ptr<arrow::RecordBatch> _batch;
+
+ // Row associated with the entry
+ row_index_t _row;
+ };
+
+ std::unordered_map<KeyType, Entry> _entries;
+
+ void store(const std::shared_ptr<RecordBatch>& batch, row_index_t row, int64_t time,
+ KeyType key) {
+ auto& e = _entries[key];
+ // that we can do this assignment optionally, is why we
+ // can get array 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;
+ }
+
+ util::optional<const Entry*> get_entry_for_key(KeyType key) const {
+ auto e = _entries.find(key);
+ if (_entries.end() == e) return util::nullopt;
+ return util::optional<const Entry*>(&e->second);
+ }
+
+ void remove_entries_with_lesser_time(int64_t ts) {
+ size_t dbg_size0 = _entries.size();
+ for (auto e = _entries.begin(); e != _entries.end();)
+ if (e->second._time < ts)
+ e = _entries.erase(e);
+ else
+ ++e;
+ size_t dbg_size1 = _entries.size();
+ if (dbg_size1 < dbg_size0) {
+ // cerr << "Removed " << dbg_size0-dbg_size1 << " memo entries.\n";
+ }
+ }
+};
+
+class InputState {
+ // InputState correponds to an input
+ // Input record batches are queued up in InputState until processed and
+ // turned into output record batches.
+
+ public:
+ InputState(const std::shared_ptr<arrow::Schema>& schema,
+ const std::string& time_col_name, const std::string& key_col_name,
+ util::optional<KeyType> wildcard_key)
+ : queue_(),
+ wildcard_key_(wildcard_key),
+ schema_(schema),
+ time_col_index_(
+ schema->GetFieldIndex(time_col_name)), // TODO: handle missing field name
+ key_col_index_(schema->GetFieldIndex(key_col_name)) {}
+
+ col_index_t init_src_to_dst_mapping(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 && is_time_or_key_column(i)))
+ src_to_dst_[i] = dst_offset++;
+ return dst_offset;
+ }
+
+ const util::optional<col_index_t>& map_src_to_dst(col_index_t src) const {
+ return src_to_dst_[src];
+ }
+
+ bool is_time_or_key_column(col_index_t i) const {
+ assert(i < schema_->num_fields());
+ return (i == time_col_index_) || (i == key_col_index_);
+ }
+
+ // Gets the latest row index, assuming the queue isn't empty
+ row_index_t get_latest_row() const { return latest_ref_row_; }
+
+ bool empty() const {
+ if (latest_ref_row_ > 0)
+ return false; // cannot be empty if ref row is >0 -- can avoid slow queue lock
+ // below
+ return queue_.empty();
+ }
+
+ int countbatches_processed_() const { return batches_processed_; }
+ int count_total_batches() const { return total_batches_; }
+
+ // Gets latest batch (precondition: must not be empty)
+ const std::shared_ptr<arrow::RecordBatch>& get_latest_batch() const {
+ return queue_.unsync_front();
+ }
+ KeyType get_latest_key() const {
+ return queue_.unsync_front()
+ ->column_data(key_col_index_)
+ ->GetValues<KeyType>(1)[latest_ref_row_];
+ }
+ int64_t get_latest_time() const {
+ return queue_.unsync_front()
+ ->column_data(time_col_index_)
+ ->GetValues<int64_t>(1)[latest_ref_row_];
+ }
+
+ bool finished() const { return batches_processed_ == total_batches_; }
+
+ bool advance() {
+ // 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();
+ if (have_active_batch) {
+ // If we have an active batch
+ if (++latest_ref_row_ >= (row_index_t)queue_.unsync_front()->num_rows()) {
+ // hit the end of the batch, need to get the next batch if possible.
+ ++batches_processed_;
+ latest_ref_row_ = 0;
+ have_active_batch &= !queue_.try_pop();
+ if (have_active_batch)
+ assert(queue_.unsync_front()->num_rows() > 0); // empty batches disallowed
+ }
+ }
+ return have_active_batch;
+ }
+
+ // Advance the data to be immediately past the specified TS, updating latest and
+ // latest_ref_row to the latest data prior to that immediate just past Returns true if
+ // updates were made, false if not.
+ bool advance_and_memoize(int64_t ts) {
+ // Advance the right side row index until we reach the latest right row (for each key)
+ // for the given left timestamp.
+
+ // Check if already updated for TS (or if there is no latest)
+ if (empty()) return false; // can't advance if empty
+ auto latest_time = get_latest_time();
+ if (latest_time > ts) return false; // already advanced
+
+ // Not updated. Try to update and possibly advance.
+ bool updated = false;
+ do {
+ latest_time = get_latest_time();
+ // 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 <= ts) {
+ memo_.store(get_latest_batch(), latest_ref_row_, latest_time, get_latest_key());
+ } else {
+ break; // hit a future timestamp -- done updating for now
+ }
+ updated = true;
+ } while (advance());
+ return updated;
+ }
+
+ void push(const std::shared_ptr<arrow::RecordBatch>& rb) {
+ if (rb->num_rows() > 0) {
+ queue_.push(rb);
+ } else {
+ ++batches_processed_; // don't enqueue empty batches, just record as processed
+ }
+ }
+
+ util::optional<const MemoStore::Entry*> get_memo_entry_for_key(KeyType key) {
+ auto r = memo_.get_entry_for_key(key);
+ if (r.has_value()) return r;
+ if (wildcard_key_.has_value()) r = memo_.get_entry_for_key(*wildcard_key_);
+ return r;
+ }
+
+ util::optional<int64_t> get_memo_time_for_key(KeyType key) {
+ auto r = get_memo_entry_for_key(key);
+ return r.has_value() ? util::make_optional((*r)->_time) : util::nullopt;
+ }
+
+ void remove_memo_entries_with_lesser_time(int64_t ts) {
+ memo_.remove_entries_with_lesser_time(ts);
+ }
+
+ const std::shared_ptr<Schema>& get_schema() const { return schema_; }
+
+ void set_total_batches(int n) {
+ assert(n >= 0);
+ assert(total_batches_ == -1); // shouldn't be set more than once
Review Comment:
```suggestion
DCHECK_GT(n, 0);
DCHECK_EQ(total_batches_, -1); // shouldn't be set more than once
```
Hmm...I'm not entirely sure total batches == 0 is an invalid case (e.g. we might filter out all data with an overly restrictive filter)
##########
cpp/src/arrow/compute/exec/asof_join_node.cc:
##########
@@ -0,0 +1,800 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+#include <iostream>
+#include <unordered_map>
+
+#include <arrow/api.h>
+#include <arrow/compute/api.h>
+#include <arrow/util/optional.h>
+#include "arrow/compute/exec/asof_join.h"
+#include "arrow/compute/exec/exec_plan.h"
+#include "arrow/compute/exec/options.h"
+#include "arrow/compute/exec/schema_util.h"
+#include "arrow/compute/exec/util.h"
+#include "arrow/result.h"
+#include "arrow/status.h"
+#include "arrow/util/checked_cast.h"
+#include "arrow/util/future.h"
+#include "arrow/util/make_unique.h"
+
+#include <condition_variable>
+#include <mutex>
+#include <thread>
+
+namespace arrow {
+namespace compute {
+
+/**
+ * Simple implementation for an unbound concurrent queue
+ */
+template <class T>
+class ConcurrentQueue {
+ public:
+ T pop() {
+ std::unique_lock<std::mutex> lock(mutex_);
+ cond_.wait(lock, [&] { return !queue_.empty(); });
+ auto item = queue_.front();
+ queue_.pop();
+ return item;
+ }
+
+ void push(const T& item) {
+ std::unique_lock<std::mutex> lock(mutex_);
+ queue_.push(item);
+ cond_.notify_one();
+ }
+
+ util::optional<T> try_pop() {
+ // Try to pop the oldest value from the queue (or return nullopt if none)
+ std::unique_lock<std::mutex> lock(mutex_);
+ if (queue_.empty()) {
+ return util::nullopt;
+ } else {
+ auto item = queue_.front();
+ queue_.pop();
+ return item;
+ }
+ }
+
+ bool empty() const {
+ std::unique_lock<std::mutex> lock(mutex_);
+ return queue_.empty();
+ }
+
+ // Un-synchronized access to front
+ // For this to be "safe":
+ // 1) the caller logically guarantees that queue is not empty
+ // 2) pop/try_pop cannot be called concurrently with this
+ const T& unsync_front() const { return queue_.front(); }
+
+ private:
+ std::queue<T> queue_;
+ mutable std::mutex mutex_;
+ std::condition_variable cond_;
+};
+
+struct MemoStore {
+ // Stores last known values for all the keys
+
+ struct Entry {
+ // Timestamp associated with the entry
+ int64_t _time;
+
+ // Batch associated with the entry (perf is probably OK for this; batches change
+ // rarely)
+ std::shared_ptr<arrow::RecordBatch> _batch;
+
+ // Row associated with the entry
+ row_index_t _row;
+ };
+
+ std::unordered_map<KeyType, Entry> _entries;
+
+ void store(const std::shared_ptr<RecordBatch>& batch, row_index_t row, int64_t time,
+ KeyType key) {
+ auto& e = _entries[key];
+ // that we can do this assignment optionally, is why we
+ // can get array 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;
+ }
+
+ util::optional<const Entry*> get_entry_for_key(KeyType key) const {
+ auto e = _entries.find(key);
+ if (_entries.end() == e) return util::nullopt;
+ return util::optional<const Entry*>(&e->second);
+ }
+
+ void remove_entries_with_lesser_time(int64_t ts) {
+ size_t dbg_size0 = _entries.size();
+ for (auto e = _entries.begin(); e != _entries.end();)
+ if (e->second._time < ts)
+ e = _entries.erase(e);
+ else
+ ++e;
+ size_t dbg_size1 = _entries.size();
+ if (dbg_size1 < dbg_size0) {
+ // cerr << "Removed " << dbg_size0-dbg_size1 << " memo entries.\n";
+ }
+ }
+};
+
+class InputState {
+ // InputState correponds to an input
+ // Input record batches are queued up in InputState until processed and
+ // turned into output record batches.
+
+ public:
+ InputState(const std::shared_ptr<arrow::Schema>& schema,
+ const std::string& time_col_name, const std::string& key_col_name,
+ util::optional<KeyType> wildcard_key)
+ : queue_(),
+ wildcard_key_(wildcard_key),
+ schema_(schema),
+ time_col_index_(
+ schema->GetFieldIndex(time_col_name)), // TODO: handle missing field name
+ key_col_index_(schema->GetFieldIndex(key_col_name)) {}
+
+ col_index_t init_src_to_dst_mapping(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 && is_time_or_key_column(i)))
+ src_to_dst_[i] = dst_offset++;
+ return dst_offset;
+ }
+
+ const util::optional<col_index_t>& map_src_to_dst(col_index_t src) const {
+ return src_to_dst_[src];
+ }
+
+ bool is_time_or_key_column(col_index_t i) const {
+ assert(i < schema_->num_fields());
+ return (i == time_col_index_) || (i == key_col_index_);
+ }
+
+ // Gets the latest row index, assuming the queue isn't empty
+ row_index_t get_latest_row() const { return latest_ref_row_; }
+
+ bool empty() const {
+ if (latest_ref_row_ > 0)
+ return false; // cannot be empty if ref row is >0 -- can avoid slow queue lock
+ // below
+ return queue_.empty();
+ }
+
+ int countbatches_processed_() const { return batches_processed_; }
+ int count_total_batches() const { return total_batches_; }
+
+ // Gets latest batch (precondition: must not be empty)
+ const std::shared_ptr<arrow::RecordBatch>& get_latest_batch() const {
+ return queue_.unsync_front();
+ }
+ KeyType get_latest_key() const {
+ return queue_.unsync_front()
+ ->column_data(key_col_index_)
+ ->GetValues<KeyType>(1)[latest_ref_row_];
+ }
+ int64_t get_latest_time() const {
+ return queue_.unsync_front()
+ ->column_data(time_col_index_)
+ ->GetValues<int64_t>(1)[latest_ref_row_];
+ }
+
+ bool finished() const { return batches_processed_ == total_batches_; }
+
+ bool advance() {
+ // 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();
+ if (have_active_batch) {
+ // If we have an active batch
+ if (++latest_ref_row_ >= (row_index_t)queue_.unsync_front()->num_rows()) {
+ // hit the end of the batch, need to get the next batch if possible.
+ ++batches_processed_;
+ latest_ref_row_ = 0;
+ have_active_batch &= !queue_.try_pop();
+ if (have_active_batch)
+ assert(queue_.unsync_front()->num_rows() > 0); // empty batches disallowed
Review Comment:
```suggestion
if (have_active_batch)
DCHECK_GT(queue_.unsync_front()->num_rows(), 0); // empty batches disallowed
```
When we improve the scheduler we can probably remove empty batches entirely. For now this is fine (I see you filter out empty batches further down)
##########
cpp/src/arrow/compute/exec/asof_join_node.cc:
##########
@@ -0,0 +1,800 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+#include <iostream>
+#include <unordered_map>
+
+#include <arrow/api.h>
+#include <arrow/compute/api.h>
+#include <arrow/util/optional.h>
+#include "arrow/compute/exec/asof_join.h"
+#include "arrow/compute/exec/exec_plan.h"
+#include "arrow/compute/exec/options.h"
+#include "arrow/compute/exec/schema_util.h"
+#include "arrow/compute/exec/util.h"
+#include "arrow/result.h"
+#include "arrow/status.h"
+#include "arrow/util/checked_cast.h"
+#include "arrow/util/future.h"
+#include "arrow/util/make_unique.h"
+
+#include <condition_variable>
+#include <mutex>
+#include <thread>
+
+namespace arrow {
+namespace compute {
+
+/**
+ * Simple implementation for an unbound concurrent queue
+ */
+template <class T>
+class ConcurrentQueue {
+ public:
+ T pop() {
+ std::unique_lock<std::mutex> lock(mutex_);
+ cond_.wait(lock, [&] { return !queue_.empty(); });
+ auto item = queue_.front();
+ queue_.pop();
+ return item;
+ }
+
+ void push(const T& item) {
+ std::unique_lock<std::mutex> lock(mutex_);
+ queue_.push(item);
+ cond_.notify_one();
+ }
+
+ util::optional<T> try_pop() {
+ // Try to pop the oldest value from the queue (or return nullopt if none)
+ std::unique_lock<std::mutex> lock(mutex_);
+ if (queue_.empty()) {
+ return util::nullopt;
+ } else {
+ auto item = queue_.front();
+ queue_.pop();
+ return item;
+ }
+ }
+
+ bool empty() const {
+ std::unique_lock<std::mutex> lock(mutex_);
+ return queue_.empty();
+ }
+
+ // Un-synchronized access to front
+ // For this to be "safe":
+ // 1) the caller logically guarantees that queue is not empty
+ // 2) pop/try_pop cannot be called concurrently with this
+ const T& unsync_front() const { return queue_.front(); }
+
+ private:
+ std::queue<T> queue_;
+ mutable std::mutex mutex_;
+ std::condition_variable cond_;
+};
+
+struct MemoStore {
+ // Stores last known values for all the keys
+
+ struct Entry {
+ // Timestamp associated with the entry
+ int64_t _time;
+
+ // Batch associated with the entry (perf is probably OK for this; batches change
+ // rarely)
+ std::shared_ptr<arrow::RecordBatch> _batch;
+
+ // Row associated with the entry
+ row_index_t _row;
+ };
+
+ std::unordered_map<KeyType, Entry> _entries;
+
+ void store(const std::shared_ptr<RecordBatch>& batch, row_index_t row, int64_t time,
+ KeyType key) {
+ auto& e = _entries[key];
+ // that we can do this assignment optionally, is why we
+ // can get array 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;
+ }
+
+ util::optional<const Entry*> get_entry_for_key(KeyType key) const {
+ auto e = _entries.find(key);
+ if (_entries.end() == e) return util::nullopt;
+ return util::optional<const Entry*>(&e->second);
+ }
+
+ void remove_entries_with_lesser_time(int64_t ts) {
+ size_t dbg_size0 = _entries.size();
+ for (auto e = _entries.begin(); e != _entries.end();)
+ if (e->second._time < ts)
+ e = _entries.erase(e);
+ else
+ ++e;
+ size_t dbg_size1 = _entries.size();
+ if (dbg_size1 < dbg_size0) {
+ // cerr << "Removed " << dbg_size0-dbg_size1 << " memo entries.\n";
+ }
+ }
+};
+
+class InputState {
+ // InputState correponds to an input
+ // Input record batches are queued up in InputState until processed and
+ // turned into output record batches.
+
+ public:
+ InputState(const std::shared_ptr<arrow::Schema>& schema,
+ const std::string& time_col_name, const std::string& key_col_name,
+ util::optional<KeyType> wildcard_key)
+ : queue_(),
+ wildcard_key_(wildcard_key),
+ schema_(schema),
+ time_col_index_(
+ schema->GetFieldIndex(time_col_name)), // TODO: handle missing field name
+ key_col_index_(schema->GetFieldIndex(key_col_name)) {}
+
+ col_index_t init_src_to_dst_mapping(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 && is_time_or_key_column(i)))
+ src_to_dst_[i] = dst_offset++;
+ return dst_offset;
+ }
+
+ const util::optional<col_index_t>& map_src_to_dst(col_index_t src) const {
+ return src_to_dst_[src];
+ }
+
+ bool is_time_or_key_column(col_index_t i) const {
+ assert(i < schema_->num_fields());
+ return (i == time_col_index_) || (i == key_col_index_);
+ }
+
+ // Gets the latest row index, assuming the queue isn't empty
+ row_index_t get_latest_row() const { return latest_ref_row_; }
+
+ bool empty() const {
+ if (latest_ref_row_ > 0)
+ return false; // cannot be empty if ref row is >0 -- can avoid slow queue lock
+ // below
+ return queue_.empty();
+ }
+
+ int countbatches_processed_() const { return batches_processed_; }
+ int count_total_batches() const { return total_batches_; }
+
+ // Gets latest batch (precondition: must not be empty)
+ const std::shared_ptr<arrow::RecordBatch>& get_latest_batch() const {
+ return queue_.unsync_front();
+ }
+ KeyType get_latest_key() const {
+ return queue_.unsync_front()
+ ->column_data(key_col_index_)
+ ->GetValues<KeyType>(1)[latest_ref_row_];
+ }
+ int64_t get_latest_time() const {
+ return queue_.unsync_front()
+ ->column_data(time_col_index_)
+ ->GetValues<int64_t>(1)[latest_ref_row_];
+ }
+
+ bool finished() const { return batches_processed_ == total_batches_; }
+
+ bool advance() {
+ // 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();
+ if (have_active_batch) {
+ // If we have an active batch
+ if (++latest_ref_row_ >= (row_index_t)queue_.unsync_front()->num_rows()) {
+ // hit the end of the batch, need to get the next batch if possible.
+ ++batches_processed_;
+ latest_ref_row_ = 0;
+ have_active_batch &= !queue_.try_pop();
+ if (have_active_batch)
+ assert(queue_.unsync_front()->num_rows() > 0); // empty batches disallowed
+ }
+ }
+ return have_active_batch;
+ }
+
+ // Advance the data to be immediately past the specified TS, updating latest and
+ // latest_ref_row to the latest data prior to that immediate just past Returns true if
+ // updates were made, false if not.
+ bool advance_and_memoize(int64_t ts) {
+ // Advance the right side row index until we reach the latest right row (for each key)
+ // for the given left timestamp.
+
+ // Check if already updated for TS (or if there is no latest)
+ if (empty()) return false; // can't advance if empty
+ auto latest_time = get_latest_time();
+ if (latest_time > ts) return false; // already advanced
+
+ // Not updated. Try to update and possibly advance.
+ bool updated = false;
+ do {
+ latest_time = get_latest_time();
+ // 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 <= ts) {
+ memo_.store(get_latest_batch(), latest_ref_row_, latest_time, get_latest_key());
+ } else {
+ break; // hit a future timestamp -- done updating for now
+ }
+ updated = true;
+ } while (advance());
+ return updated;
+ }
+
+ void push(const std::shared_ptr<arrow::RecordBatch>& rb) {
+ if (rb->num_rows() > 0) {
+ queue_.push(rb);
+ } else {
+ ++batches_processed_; // don't enqueue empty batches, just record as processed
+ }
+ }
+
+ util::optional<const MemoStore::Entry*> get_memo_entry_for_key(KeyType key) {
+ auto r = memo_.get_entry_for_key(key);
+ if (r.has_value()) return r;
+ if (wildcard_key_.has_value()) r = memo_.get_entry_for_key(*wildcard_key_);
+ return r;
+ }
+
+ util::optional<int64_t> get_memo_time_for_key(KeyType key) {
+ auto r = get_memo_entry_for_key(key);
+ return r.has_value() ? util::make_optional((*r)->_time) : util::nullopt;
+ }
+
+ void remove_memo_entries_with_lesser_time(int64_t ts) {
+ memo_.remove_entries_with_lesser_time(ts);
+ }
+
+ const std::shared_ptr<Schema>& get_schema() const { return schema_; }
+
+ void set_total_batches(int n) {
+ assert(n >= 0);
+ assert(total_batches_ == -1); // shouldn't be set more than once
+ total_batches_ = n;
+ }
+
+ private:
+ // Pending record batches. The latest is the front. Batches cannot be empty.
+ ConcurrentQueue<std::shared_ptr<RecordBatch>> queue_;
+
+ // Wildcard key for this input, if applicable.
+ util::optional<KeyType> wildcard_key_;
+
+ // Schema associated with the input
+ std::shared_ptr<Schema> schema_;
+
+ // Total number of batches (only int because InputFinished uses int)
+ int total_batches_ = -1;
+
+ // Number of batches processed so far (only int because InputFinished uses int)
+ int batches_processed_ = 0;
+
+ // Index of the time col
+ col_index_t time_col_index_;
+
+ // Index of the key col
+ col_index_t key_col_index_;
+
+ // Index of the latest row reference within; if >0 then queue_ cannot be empty
+ row_index_t latest_ref_row_ =
+ 0; // must be < queue_.front()->num_rows() if queue_ is non-empty
+
+ // Stores latest known values for the various keys
+ MemoStore memo_;
+
+ // Mapping of source columns to destination columns
+ std::vector<util::optional<col_index_t>> src_to_dst_;
+};
+
+template <size_t MAX_TABLES>
+struct CompositeReferenceRow {
+ struct Entry {
+ arrow::RecordBatch* batch; // can be NULL if there's no value
+ row_index_t row;
+ };
+ Entry refs[MAX_TABLES];
+};
+
+// A table of composite reference rows. Rows maintain pointers to the
+// constituent record batches, but the overall table retains shared_ptr
+// references to ensure memory remains resident while the table is live.
+//
+// The main reason for this is that, especially for wide tables, joins
+// are effectively row-oriented, rather than column-oriented. Separating
+// the join part from the columnar materialization part simplifies the
+// logic around data types and increases efficiency.
+//
+// We don't put the shared_ptr's into the rows for efficiency reasons.
+template <size_t MAX_TABLES>
+class CompositeReferenceTable {
+ public:
+ explicit CompositeReferenceTable(size_t n_tables) : n_tables_(n_tables) {
+ assert(n_tables_ >= 1);
+ assert(n_tables_ <= MAX_TABLES);
+ }
+
+ size_t n_rows() const { return rows_.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<std::unique_ptr<InputState>>& in, int64_t tolerance) {
+ assert(in.size() == n_tables_);
Review Comment:
```suggestion
DCHECK_EQ(in.size(), n_tables_);
```
##########
cpp/src/arrow/compute/exec/asof_join_node.cc:
##########
@@ -0,0 +1,800 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+#include <iostream>
+#include <unordered_map>
+
+#include <arrow/api.h>
+#include <arrow/compute/api.h>
+#include <arrow/util/optional.h>
+#include "arrow/compute/exec/asof_join.h"
+#include "arrow/compute/exec/exec_plan.h"
+#include "arrow/compute/exec/options.h"
+#include "arrow/compute/exec/schema_util.h"
+#include "arrow/compute/exec/util.h"
+#include "arrow/result.h"
+#include "arrow/status.h"
+#include "arrow/util/checked_cast.h"
+#include "arrow/util/future.h"
+#include "arrow/util/make_unique.h"
+
+#include <condition_variable>
+#include <mutex>
+#include <thread>
+
+namespace arrow {
+namespace compute {
+
+/**
+ * Simple implementation for an unbound concurrent queue
+ */
+template <class T>
+class ConcurrentQueue {
+ public:
+ T pop() {
+ std::unique_lock<std::mutex> lock(mutex_);
+ cond_.wait(lock, [&] { return !queue_.empty(); });
+ auto item = queue_.front();
+ queue_.pop();
+ return item;
+ }
+
+ void push(const T& item) {
+ std::unique_lock<std::mutex> lock(mutex_);
+ queue_.push(item);
+ cond_.notify_one();
+ }
+
+ util::optional<T> try_pop() {
+ // Try to pop the oldest value from the queue (or return nullopt if none)
+ std::unique_lock<std::mutex> lock(mutex_);
+ if (queue_.empty()) {
+ return util::nullopt;
+ } else {
+ auto item = queue_.front();
+ queue_.pop();
+ return item;
+ }
+ }
+
+ bool empty() const {
+ std::unique_lock<std::mutex> lock(mutex_);
+ return queue_.empty();
+ }
+
+ // Un-synchronized access to front
+ // For this to be "safe":
+ // 1) the caller logically guarantees that queue is not empty
+ // 2) pop/try_pop cannot be called concurrently with this
+ const T& unsync_front() const { return queue_.front(); }
+
+ private:
+ std::queue<T> queue_;
+ mutable std::mutex mutex_;
+ std::condition_variable cond_;
+};
+
+struct MemoStore {
+ // Stores last known values for all the keys
+
+ struct Entry {
+ // Timestamp associated with the entry
+ int64_t _time;
+
+ // Batch associated with the entry (perf is probably OK for this; batches change
+ // rarely)
+ std::shared_ptr<arrow::RecordBatch> _batch;
+
+ // Row associated with the entry
+ row_index_t _row;
+ };
+
+ std::unordered_map<KeyType, Entry> _entries;
+
+ void store(const std::shared_ptr<RecordBatch>& batch, row_index_t row, int64_t time,
+ KeyType key) {
+ auto& e = _entries[key];
+ // that we can do this assignment optionally, is why we
+ // can get array 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;
+ }
+
+ util::optional<const Entry*> get_entry_for_key(KeyType key) const {
+ auto e = _entries.find(key);
+ if (_entries.end() == e) return util::nullopt;
+ return util::optional<const Entry*>(&e->second);
+ }
+
+ void remove_entries_with_lesser_time(int64_t ts) {
+ size_t dbg_size0 = _entries.size();
+ for (auto e = _entries.begin(); e != _entries.end();)
+ if (e->second._time < ts)
+ e = _entries.erase(e);
+ else
+ ++e;
+ size_t dbg_size1 = _entries.size();
+ if (dbg_size1 < dbg_size0) {
+ // cerr << "Removed " << dbg_size0-dbg_size1 << " memo entries.\n";
+ }
+ }
+};
+
+class InputState {
+ // InputState correponds to an input
+ // Input record batches are queued up in InputState until processed and
+ // turned into output record batches.
+
+ public:
+ InputState(const std::shared_ptr<arrow::Schema>& schema,
+ const std::string& time_col_name, const std::string& key_col_name,
+ util::optional<KeyType> wildcard_key)
+ : queue_(),
+ wildcard_key_(wildcard_key),
+ schema_(schema),
+ time_col_index_(
+ schema->GetFieldIndex(time_col_name)), // TODO: handle missing field name
+ key_col_index_(schema->GetFieldIndex(key_col_name)) {}
+
+ col_index_t init_src_to_dst_mapping(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 && is_time_or_key_column(i)))
+ src_to_dst_[i] = dst_offset++;
+ return dst_offset;
+ }
+
+ const util::optional<col_index_t>& map_src_to_dst(col_index_t src) const {
+ return src_to_dst_[src];
+ }
+
+ bool is_time_or_key_column(col_index_t i) const {
+ assert(i < schema_->num_fields());
+ return (i == time_col_index_) || (i == key_col_index_);
+ }
+
+ // Gets the latest row index, assuming the queue isn't empty
+ row_index_t get_latest_row() const { return latest_ref_row_; }
+
+ bool empty() const {
+ if (latest_ref_row_ > 0)
+ return false; // cannot be empty if ref row is >0 -- can avoid slow queue lock
+ // below
+ return queue_.empty();
+ }
+
+ int countbatches_processed_() const { return batches_processed_; }
+ int count_total_batches() const { return total_batches_; }
+
+ // Gets latest batch (precondition: must not be empty)
+ const std::shared_ptr<arrow::RecordBatch>& get_latest_batch() const {
+ return queue_.unsync_front();
+ }
+ KeyType get_latest_key() const {
+ return queue_.unsync_front()
+ ->column_data(key_col_index_)
+ ->GetValues<KeyType>(1)[latest_ref_row_];
+ }
+ int64_t get_latest_time() const {
+ return queue_.unsync_front()
+ ->column_data(time_col_index_)
+ ->GetValues<int64_t>(1)[latest_ref_row_];
+ }
+
+ bool finished() const { return batches_processed_ == total_batches_; }
+
+ bool advance() {
+ // 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();
+ if (have_active_batch) {
+ // If we have an active batch
+ if (++latest_ref_row_ >= (row_index_t)queue_.unsync_front()->num_rows()) {
+ // hit the end of the batch, need to get the next batch if possible.
+ ++batches_processed_;
+ latest_ref_row_ = 0;
+ have_active_batch &= !queue_.try_pop();
+ if (have_active_batch)
+ assert(queue_.unsync_front()->num_rows() > 0); // empty batches disallowed
+ }
+ }
+ return have_active_batch;
+ }
+
+ // Advance the data to be immediately past the specified TS, updating latest and
+ // latest_ref_row to the latest data prior to that immediate just past Returns true if
+ // updates were made, false if not.
+ bool advance_and_memoize(int64_t ts) {
+ // Advance the right side row index until we reach the latest right row (for each key)
+ // for the given left timestamp.
+
+ // Check if already updated for TS (or if there is no latest)
+ if (empty()) return false; // can't advance if empty
+ auto latest_time = get_latest_time();
+ if (latest_time > ts) return false; // already advanced
+
+ // Not updated. Try to update and possibly advance.
+ bool updated = false;
+ do {
+ latest_time = get_latest_time();
+ // 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 <= ts) {
+ memo_.store(get_latest_batch(), latest_ref_row_, latest_time, get_latest_key());
+ } else {
+ break; // hit a future timestamp -- done updating for now
+ }
+ updated = true;
+ } while (advance());
+ return updated;
+ }
+
+ void push(const std::shared_ptr<arrow::RecordBatch>& rb) {
+ if (rb->num_rows() > 0) {
+ queue_.push(rb);
+ } else {
+ ++batches_processed_; // don't enqueue empty batches, just record as processed
+ }
+ }
+
+ util::optional<const MemoStore::Entry*> get_memo_entry_for_key(KeyType key) {
+ auto r = memo_.get_entry_for_key(key);
+ if (r.has_value()) return r;
+ if (wildcard_key_.has_value()) r = memo_.get_entry_for_key(*wildcard_key_);
+ return r;
+ }
+
+ util::optional<int64_t> get_memo_time_for_key(KeyType key) {
+ auto r = get_memo_entry_for_key(key);
+ return r.has_value() ? util::make_optional((*r)->_time) : util::nullopt;
+ }
+
+ void remove_memo_entries_with_lesser_time(int64_t ts) {
+ memo_.remove_entries_with_lesser_time(ts);
+ }
+
+ const std::shared_ptr<Schema>& get_schema() const { return schema_; }
+
+ void set_total_batches(int n) {
+ assert(n >= 0);
+ assert(total_batches_ == -1); // shouldn't be set more than once
+ total_batches_ = n;
+ }
+
+ private:
+ // Pending record batches. The latest is the front. Batches cannot be empty.
+ ConcurrentQueue<std::shared_ptr<RecordBatch>> queue_;
+
+ // Wildcard key for this input, if applicable.
+ util::optional<KeyType> wildcard_key_;
+
+ // Schema associated with the input
+ std::shared_ptr<Schema> schema_;
+
+ // Total number of batches (only int because InputFinished uses int)
+ int total_batches_ = -1;
+
+ // Number of batches processed so far (only int because InputFinished uses int)
+ int batches_processed_ = 0;
+
+ // Index of the time col
+ col_index_t time_col_index_;
+
+ // Index of the key col
+ col_index_t key_col_index_;
+
+ // Index of the latest row reference within; if >0 then queue_ cannot be empty
+ row_index_t latest_ref_row_ =
+ 0; // must be < queue_.front()->num_rows() if queue_ is non-empty
+
+ // Stores latest known values for the various keys
+ MemoStore memo_;
+
+ // Mapping of source columns to destination columns
+ std::vector<util::optional<col_index_t>> src_to_dst_;
+};
+
+template <size_t MAX_TABLES>
+struct CompositeReferenceRow {
+ struct Entry {
+ arrow::RecordBatch* batch; // can be NULL if there's no value
+ row_index_t row;
+ };
+ Entry refs[MAX_TABLES];
+};
+
+// A table of composite reference rows. Rows maintain pointers to the
+// constituent record batches, but the overall table retains shared_ptr
+// references to ensure memory remains resident while the table is live.
+//
+// The main reason for this is that, especially for wide tables, joins
+// are effectively row-oriented, rather than column-oriented. Separating
+// the join part from the columnar materialization part simplifies the
+// logic around data types and increases efficiency.
+//
+// We don't put the shared_ptr's into the rows for efficiency reasons.
+template <size_t MAX_TABLES>
+class CompositeReferenceTable {
+ public:
+ explicit CompositeReferenceTable(size_t n_tables) : n_tables_(n_tables) {
+ assert(n_tables_ >= 1);
+ assert(n_tables_ <= MAX_TABLES);
Review Comment:
```suggestion
DCHECK_GE(n_tables_, 1);
DCHECK_LE(n_tables_, MAX_TABLES);
```
##########
cpp/src/arrow/compute/exec/asof_join_node.cc:
##########
@@ -0,0 +1,800 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+#include <iostream>
+#include <unordered_map>
+
+#include <arrow/api.h>
+#include <arrow/compute/api.h>
+#include <arrow/util/optional.h>
+#include "arrow/compute/exec/asof_join.h"
+#include "arrow/compute/exec/exec_plan.h"
+#include "arrow/compute/exec/options.h"
+#include "arrow/compute/exec/schema_util.h"
+#include "arrow/compute/exec/util.h"
+#include "arrow/result.h"
+#include "arrow/status.h"
+#include "arrow/util/checked_cast.h"
+#include "arrow/util/future.h"
+#include "arrow/util/make_unique.h"
+
+#include <condition_variable>
+#include <mutex>
+#include <thread>
+
+namespace arrow {
+namespace compute {
+
+/**
+ * Simple implementation for an unbound concurrent queue
+ */
+template <class T>
+class ConcurrentQueue {
+ public:
+ T pop() {
+ std::unique_lock<std::mutex> lock(mutex_);
+ cond_.wait(lock, [&] { return !queue_.empty(); });
+ auto item = queue_.front();
+ queue_.pop();
+ return item;
+ }
+
+ void push(const T& item) {
+ std::unique_lock<std::mutex> lock(mutex_);
+ queue_.push(item);
+ cond_.notify_one();
+ }
+
+ util::optional<T> try_pop() {
+ // Try to pop the oldest value from the queue (or return nullopt if none)
+ std::unique_lock<std::mutex> lock(mutex_);
+ if (queue_.empty()) {
+ return util::nullopt;
+ } else {
+ auto item = queue_.front();
+ queue_.pop();
+ return item;
+ }
+ }
+
+ bool empty() const {
+ std::unique_lock<std::mutex> lock(mutex_);
+ return queue_.empty();
+ }
+
+ // Un-synchronized access to front
+ // For this to be "safe":
+ // 1) the caller logically guarantees that queue is not empty
+ // 2) pop/try_pop cannot be called concurrently with this
+ const T& unsync_front() const { return queue_.front(); }
+
+ private:
+ std::queue<T> queue_;
+ mutable std::mutex mutex_;
+ std::condition_variable cond_;
+};
+
+struct MemoStore {
+ // Stores last known values for all the keys
+
+ struct Entry {
+ // Timestamp associated with the entry
+ int64_t _time;
+
+ // Batch associated with the entry (perf is probably OK for this; batches change
+ // rarely)
+ std::shared_ptr<arrow::RecordBatch> _batch;
+
+ // Row associated with the entry
+ row_index_t _row;
+ };
+
+ std::unordered_map<KeyType, Entry> _entries;
+
+ void store(const std::shared_ptr<RecordBatch>& batch, row_index_t row, int64_t time,
+ KeyType key) {
+ auto& e = _entries[key];
+ // that we can do this assignment optionally, is why we
+ // can get array 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;
+ }
+
+ util::optional<const Entry*> get_entry_for_key(KeyType key) const {
+ auto e = _entries.find(key);
+ if (_entries.end() == e) return util::nullopt;
+ return util::optional<const Entry*>(&e->second);
+ }
+
+ void remove_entries_with_lesser_time(int64_t ts) {
+ size_t dbg_size0 = _entries.size();
+ for (auto e = _entries.begin(); e != _entries.end();)
+ if (e->second._time < ts)
+ e = _entries.erase(e);
+ else
+ ++e;
+ size_t dbg_size1 = _entries.size();
+ if (dbg_size1 < dbg_size0) {
+ // cerr << "Removed " << dbg_size0-dbg_size1 << " memo entries.\n";
+ }
+ }
+};
+
+class InputState {
+ // InputState correponds to an input
+ // Input record batches are queued up in InputState until processed and
+ // turned into output record batches.
+
+ public:
+ InputState(const std::shared_ptr<arrow::Schema>& schema,
+ const std::string& time_col_name, const std::string& key_col_name,
+ util::optional<KeyType> wildcard_key)
+ : queue_(),
+ wildcard_key_(wildcard_key),
+ schema_(schema),
+ time_col_index_(
+ schema->GetFieldIndex(time_col_name)), // TODO: handle missing field name
+ key_col_index_(schema->GetFieldIndex(key_col_name)) {}
+
+ col_index_t init_src_to_dst_mapping(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 && is_time_or_key_column(i)))
+ src_to_dst_[i] = dst_offset++;
+ return dst_offset;
+ }
+
+ const util::optional<col_index_t>& map_src_to_dst(col_index_t src) const {
+ return src_to_dst_[src];
+ }
+
+ bool is_time_or_key_column(col_index_t i) const {
+ assert(i < schema_->num_fields());
+ return (i == time_col_index_) || (i == key_col_index_);
+ }
+
+ // Gets the latest row index, assuming the queue isn't empty
+ row_index_t get_latest_row() const { return latest_ref_row_; }
+
+ bool empty() const {
+ if (latest_ref_row_ > 0)
+ return false; // cannot be empty if ref row is >0 -- can avoid slow queue lock
+ // below
+ return queue_.empty();
+ }
+
+ int countbatches_processed_() const { return batches_processed_; }
+ int count_total_batches() const { return total_batches_; }
+
+ // Gets latest batch (precondition: must not be empty)
+ const std::shared_ptr<arrow::RecordBatch>& get_latest_batch() const {
+ return queue_.unsync_front();
+ }
+ KeyType get_latest_key() const {
+ return queue_.unsync_front()
+ ->column_data(key_col_index_)
+ ->GetValues<KeyType>(1)[latest_ref_row_];
+ }
+ int64_t get_latest_time() const {
+ return queue_.unsync_front()
+ ->column_data(time_col_index_)
+ ->GetValues<int64_t>(1)[latest_ref_row_];
+ }
+
+ bool finished() const { return batches_processed_ == total_batches_; }
+
+ bool advance() {
+ // 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();
+ if (have_active_batch) {
+ // If we have an active batch
+ if (++latest_ref_row_ >= (row_index_t)queue_.unsync_front()->num_rows()) {
+ // hit the end of the batch, need to get the next batch if possible.
+ ++batches_processed_;
+ latest_ref_row_ = 0;
+ have_active_batch &= !queue_.try_pop();
+ if (have_active_batch)
+ assert(queue_.unsync_front()->num_rows() > 0); // empty batches disallowed
+ }
+ }
+ return have_active_batch;
+ }
+
+ // Advance the data to be immediately past the specified TS, updating latest and
+ // latest_ref_row to the latest data prior to that immediate just past Returns true if
+ // updates were made, false if not.
+ bool advance_and_memoize(int64_t ts) {
+ // Advance the right side row index until we reach the latest right row (for each key)
+ // for the given left timestamp.
+
+ // Check if already updated for TS (or if there is no latest)
+ if (empty()) return false; // can't advance if empty
+ auto latest_time = get_latest_time();
+ if (latest_time > ts) return false; // already advanced
+
+ // Not updated. Try to update and possibly advance.
+ bool updated = false;
+ do {
+ latest_time = get_latest_time();
+ // 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 <= ts) {
+ memo_.store(get_latest_batch(), latest_ref_row_, latest_time, get_latest_key());
+ } else {
+ break; // hit a future timestamp -- done updating for now
+ }
+ updated = true;
+ } while (advance());
+ return updated;
+ }
+
+ void push(const std::shared_ptr<arrow::RecordBatch>& rb) {
+ if (rb->num_rows() > 0) {
+ queue_.push(rb);
+ } else {
+ ++batches_processed_; // don't enqueue empty batches, just record as processed
+ }
+ }
+
+ util::optional<const MemoStore::Entry*> get_memo_entry_for_key(KeyType key) {
Review Comment:
If you are returning a pointer here why not use `nullptr` as your "no value found" sentinel? Then you don't need `util::optional`
##########
cpp/src/arrow/compute/exec/asof_join_node.cc:
##########
@@ -0,0 +1,800 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+#include <iostream>
+#include <unordered_map>
+
+#include <arrow/api.h>
+#include <arrow/compute/api.h>
+#include <arrow/util/optional.h>
+#include "arrow/compute/exec/asof_join.h"
+#include "arrow/compute/exec/exec_plan.h"
+#include "arrow/compute/exec/options.h"
+#include "arrow/compute/exec/schema_util.h"
+#include "arrow/compute/exec/util.h"
+#include "arrow/result.h"
+#include "arrow/status.h"
+#include "arrow/util/checked_cast.h"
+#include "arrow/util/future.h"
+#include "arrow/util/make_unique.h"
+
+#include <condition_variable>
+#include <mutex>
+#include <thread>
+
+namespace arrow {
+namespace compute {
+
+/**
+ * Simple implementation for an unbound concurrent queue
+ */
+template <class T>
+class ConcurrentQueue {
+ public:
+ T pop() {
+ std::unique_lock<std::mutex> lock(mutex_);
+ cond_.wait(lock, [&] { return !queue_.empty(); });
+ auto item = queue_.front();
+ queue_.pop();
+ return item;
+ }
+
+ void push(const T& item) {
+ std::unique_lock<std::mutex> lock(mutex_);
+ queue_.push(item);
+ cond_.notify_one();
+ }
+
+ util::optional<T> try_pop() {
+ // Try to pop the oldest value from the queue (or return nullopt if none)
+ std::unique_lock<std::mutex> lock(mutex_);
+ if (queue_.empty()) {
+ return util::nullopt;
+ } else {
+ auto item = queue_.front();
+ queue_.pop();
+ return item;
+ }
+ }
+
+ bool empty() const {
+ std::unique_lock<std::mutex> lock(mutex_);
+ return queue_.empty();
+ }
+
+ // Un-synchronized access to front
+ // For this to be "safe":
+ // 1) the caller logically guarantees that queue is not empty
+ // 2) pop/try_pop cannot be called concurrently with this
+ const T& unsync_front() const { return queue_.front(); }
+
+ private:
+ std::queue<T> queue_;
+ mutable std::mutex mutex_;
+ std::condition_variable cond_;
+};
+
+struct MemoStore {
+ // Stores last known values for all the keys
+
+ struct Entry {
+ // Timestamp associated with the entry
+ int64_t _time;
+
+ // Batch associated with the entry (perf is probably OK for this; batches change
+ // rarely)
+ std::shared_ptr<arrow::RecordBatch> _batch;
+
+ // Row associated with the entry
+ row_index_t _row;
+ };
+
+ std::unordered_map<KeyType, Entry> _entries;
+
+ void store(const std::shared_ptr<RecordBatch>& batch, row_index_t row, int64_t time,
+ KeyType key) {
+ auto& e = _entries[key];
+ // that we can do this assignment optionally, is why we
+ // can get array 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;
+ }
+
+ util::optional<const Entry*> get_entry_for_key(KeyType key) const {
+ auto e = _entries.find(key);
+ if (_entries.end() == e) return util::nullopt;
+ return util::optional<const Entry*>(&e->second);
+ }
+
+ void remove_entries_with_lesser_time(int64_t ts) {
+ size_t dbg_size0 = _entries.size();
+ for (auto e = _entries.begin(); e != _entries.end();)
+ if (e->second._time < ts)
+ e = _entries.erase(e);
+ else
+ ++e;
+ size_t dbg_size1 = _entries.size();
+ if (dbg_size1 < dbg_size0) {
+ // cerr << "Removed " << dbg_size0-dbg_size1 << " memo entries.\n";
+ }
+ }
+};
+
+class InputState {
+ // InputState correponds to an input
+ // Input record batches are queued up in InputState until processed and
+ // turned into output record batches.
+
+ public:
+ InputState(const std::shared_ptr<arrow::Schema>& schema,
+ const std::string& time_col_name, const std::string& key_col_name,
+ util::optional<KeyType> wildcard_key)
+ : queue_(),
+ wildcard_key_(wildcard_key),
+ schema_(schema),
+ time_col_index_(
+ schema->GetFieldIndex(time_col_name)), // TODO: handle missing field name
+ key_col_index_(schema->GetFieldIndex(key_col_name)) {}
+
+ col_index_t init_src_to_dst_mapping(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 && is_time_or_key_column(i)))
+ src_to_dst_[i] = dst_offset++;
+ return dst_offset;
+ }
+
+ const util::optional<col_index_t>& map_src_to_dst(col_index_t src) const {
+ return src_to_dst_[src];
+ }
+
+ bool is_time_or_key_column(col_index_t i) const {
+ assert(i < schema_->num_fields());
+ return (i == time_col_index_) || (i == key_col_index_);
+ }
+
+ // Gets the latest row index, assuming the queue isn't empty
+ row_index_t get_latest_row() const { return latest_ref_row_; }
+
+ bool empty() const {
+ if (latest_ref_row_ > 0)
+ return false; // cannot be empty if ref row is >0 -- can avoid slow queue lock
+ // below
+ return queue_.empty();
+ }
+
+ int countbatches_processed_() const { return batches_processed_; }
+ int count_total_batches() const { return total_batches_; }
+
+ // Gets latest batch (precondition: must not be empty)
+ const std::shared_ptr<arrow::RecordBatch>& get_latest_batch() const {
+ return queue_.unsync_front();
+ }
+ KeyType get_latest_key() const {
+ return queue_.unsync_front()
+ ->column_data(key_col_index_)
+ ->GetValues<KeyType>(1)[latest_ref_row_];
+ }
+ int64_t get_latest_time() const {
+ return queue_.unsync_front()
+ ->column_data(time_col_index_)
+ ->GetValues<int64_t>(1)[latest_ref_row_];
+ }
+
+ bool finished() const { return batches_processed_ == total_batches_; }
+
+ bool advance() {
+ // 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();
+ if (have_active_batch) {
+ // If we have an active batch
+ if (++latest_ref_row_ >= (row_index_t)queue_.unsync_front()->num_rows()) {
+ // hit the end of the batch, need to get the next batch if possible.
+ ++batches_processed_;
+ latest_ref_row_ = 0;
+ have_active_batch &= !queue_.try_pop();
+ if (have_active_batch)
+ assert(queue_.unsync_front()->num_rows() > 0); // empty batches disallowed
+ }
+ }
+ return have_active_batch;
+ }
+
+ // Advance the data to be immediately past the specified TS, updating latest and
+ // latest_ref_row to the latest data prior to that immediate just past Returns true if
+ // updates were made, false if not.
+ bool advance_and_memoize(int64_t ts) {
+ // Advance the right side row index until we reach the latest right row (for each key)
+ // for the given left timestamp.
+
+ // Check if already updated for TS (or if there is no latest)
+ if (empty()) return false; // can't advance if empty
+ auto latest_time = get_latest_time();
+ if (latest_time > ts) return false; // already advanced
+
+ // Not updated. Try to update and possibly advance.
+ bool updated = false;
+ do {
+ latest_time = get_latest_time();
+ // 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 <= ts) {
+ memo_.store(get_latest_batch(), latest_ref_row_, latest_time, get_latest_key());
+ } else {
+ break; // hit a future timestamp -- done updating for now
+ }
+ updated = true;
+ } while (advance());
+ return updated;
+ }
+
+ void push(const std::shared_ptr<arrow::RecordBatch>& rb) {
+ if (rb->num_rows() > 0) {
+ queue_.push(rb);
+ } else {
+ ++batches_processed_; // don't enqueue empty batches, just record as processed
+ }
+ }
+
+ util::optional<const MemoStore::Entry*> get_memo_entry_for_key(KeyType key) {
+ auto r = memo_.get_entry_for_key(key);
+ if (r.has_value()) return r;
+ if (wildcard_key_.has_value()) r = memo_.get_entry_for_key(*wildcard_key_);
+ return r;
+ }
+
+ util::optional<int64_t> get_memo_time_for_key(KeyType key) {
+ auto r = get_memo_entry_for_key(key);
+ return r.has_value() ? util::make_optional((*r)->_time) : util::nullopt;
+ }
+
+ void remove_memo_entries_with_lesser_time(int64_t ts) {
+ memo_.remove_entries_with_lesser_time(ts);
+ }
+
+ const std::shared_ptr<Schema>& get_schema() const { return schema_; }
+
+ void set_total_batches(int n) {
+ assert(n >= 0);
+ assert(total_batches_ == -1); // shouldn't be set more than once
+ total_batches_ = n;
+ }
+
+ private:
+ // Pending record batches. The latest is the front. Batches cannot be empty.
+ ConcurrentQueue<std::shared_ptr<RecordBatch>> queue_;
+
+ // Wildcard key for this input, if applicable.
+ util::optional<KeyType> wildcard_key_;
+
+ // Schema associated with the input
+ std::shared_ptr<Schema> schema_;
+
+ // Total number of batches (only int because InputFinished uses int)
+ int total_batches_ = -1;
+
+ // Number of batches processed so far (only int because InputFinished uses int)
+ int batches_processed_ = 0;
+
+ // Index of the time col
+ col_index_t time_col_index_;
+
+ // Index of the key col
+ col_index_t key_col_index_;
+
+ // Index of the latest row reference within; if >0 then queue_ cannot be empty
+ row_index_t latest_ref_row_ =
+ 0; // must be < queue_.front()->num_rows() if queue_ is non-empty
+
+ // Stores latest known values for the various keys
+ MemoStore memo_;
+
+ // Mapping of source columns to destination columns
+ std::vector<util::optional<col_index_t>> src_to_dst_;
+};
+
+template <size_t MAX_TABLES>
+struct CompositeReferenceRow {
+ struct Entry {
+ arrow::RecordBatch* batch; // can be NULL if there's no value
+ row_index_t row;
+ };
+ Entry refs[MAX_TABLES];
Review Comment:
Why not use `std::vector` here and avoid the template parameter? There should only be one allocation per plan which is pretty minor.
##########
cpp/src/arrow/compute/exec/asof_join_node.cc:
##########
@@ -0,0 +1,800 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+#include <iostream>
+#include <unordered_map>
+
+#include <arrow/api.h>
+#include <arrow/compute/api.h>
+#include <arrow/util/optional.h>
+#include "arrow/compute/exec/asof_join.h"
+#include "arrow/compute/exec/exec_plan.h"
+#include "arrow/compute/exec/options.h"
+#include "arrow/compute/exec/schema_util.h"
+#include "arrow/compute/exec/util.h"
+#include "arrow/result.h"
+#include "arrow/status.h"
+#include "arrow/util/checked_cast.h"
+#include "arrow/util/future.h"
+#include "arrow/util/make_unique.h"
+
+#include <condition_variable>
+#include <mutex>
+#include <thread>
+
+namespace arrow {
+namespace compute {
+
+/**
+ * Simple implementation for an unbound concurrent queue
+ */
+template <class T>
+class ConcurrentQueue {
+ public:
+ T pop() {
+ std::unique_lock<std::mutex> lock(mutex_);
+ cond_.wait(lock, [&] { return !queue_.empty(); });
+ auto item = queue_.front();
+ queue_.pop();
+ return item;
+ }
+
+ void push(const T& item) {
+ std::unique_lock<std::mutex> lock(mutex_);
+ queue_.push(item);
+ cond_.notify_one();
+ }
+
+ util::optional<T> try_pop() {
+ // Try to pop the oldest value from the queue (or return nullopt if none)
+ std::unique_lock<std::mutex> lock(mutex_);
+ if (queue_.empty()) {
+ return util::nullopt;
+ } else {
+ auto item = queue_.front();
+ queue_.pop();
+ return item;
+ }
+ }
+
+ bool empty() const {
+ std::unique_lock<std::mutex> lock(mutex_);
+ return queue_.empty();
+ }
+
+ // Un-synchronized access to front
+ // For this to be "safe":
+ // 1) the caller logically guarantees that queue is not empty
+ // 2) pop/try_pop cannot be called concurrently with this
+ const T& unsync_front() const { return queue_.front(); }
+
+ private:
+ std::queue<T> queue_;
+ mutable std::mutex mutex_;
+ std::condition_variable cond_;
+};
+
+struct MemoStore {
+ // Stores last known values for all the keys
+
+ struct Entry {
+ // Timestamp associated with the entry
+ int64_t _time;
+
+ // Batch associated with the entry (perf is probably OK for this; batches change
+ // rarely)
+ std::shared_ptr<arrow::RecordBatch> _batch;
+
+ // Row associated with the entry
+ row_index_t _row;
+ };
+
+ std::unordered_map<KeyType, Entry> _entries;
+
+ void store(const std::shared_ptr<RecordBatch>& batch, row_index_t row, int64_t time,
+ KeyType key) {
+ auto& e = _entries[key];
+ // that we can do this assignment optionally, is why we
+ // can get array 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;
+ }
+
+ util::optional<const Entry*> get_entry_for_key(KeyType key) const {
+ auto e = _entries.find(key);
+ if (_entries.end() == e) return util::nullopt;
+ return util::optional<const Entry*>(&e->second);
+ }
+
+ void remove_entries_with_lesser_time(int64_t ts) {
+ size_t dbg_size0 = _entries.size();
+ for (auto e = _entries.begin(); e != _entries.end();)
+ if (e->second._time < ts)
+ e = _entries.erase(e);
+ else
+ ++e;
+ size_t dbg_size1 = _entries.size();
+ if (dbg_size1 < dbg_size0) {
+ // cerr << "Removed " << dbg_size0-dbg_size1 << " memo entries.\n";
+ }
+ }
+};
+
+class InputState {
+ // InputState correponds to an input
+ // Input record batches are queued up in InputState until processed and
+ // turned into output record batches.
+
+ public:
+ InputState(const std::shared_ptr<arrow::Schema>& schema,
+ const std::string& time_col_name, const std::string& key_col_name,
+ util::optional<KeyType> wildcard_key)
+ : queue_(),
+ wildcard_key_(wildcard_key),
+ schema_(schema),
+ time_col_index_(
+ schema->GetFieldIndex(time_col_name)), // TODO: handle missing field name
+ key_col_index_(schema->GetFieldIndex(key_col_name)) {}
+
+ col_index_t init_src_to_dst_mapping(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 && is_time_or_key_column(i)))
+ src_to_dst_[i] = dst_offset++;
+ return dst_offset;
+ }
+
+ const util::optional<col_index_t>& map_src_to_dst(col_index_t src) const {
+ return src_to_dst_[src];
+ }
+
+ bool is_time_or_key_column(col_index_t i) const {
+ assert(i < schema_->num_fields());
+ return (i == time_col_index_) || (i == key_col_index_);
+ }
+
+ // Gets the latest row index, assuming the queue isn't empty
+ row_index_t get_latest_row() const { return latest_ref_row_; }
+
+ bool empty() const {
+ if (latest_ref_row_ > 0)
+ return false; // cannot be empty if ref row is >0 -- can avoid slow queue lock
+ // below
+ return queue_.empty();
+ }
+
+ int countbatches_processed_() const { return batches_processed_; }
+ int count_total_batches() const { return total_batches_; }
+
+ // Gets latest batch (precondition: must not be empty)
+ const std::shared_ptr<arrow::RecordBatch>& get_latest_batch() const {
+ return queue_.unsync_front();
+ }
+ KeyType get_latest_key() const {
+ return queue_.unsync_front()
+ ->column_data(key_col_index_)
+ ->GetValues<KeyType>(1)[latest_ref_row_];
+ }
+ int64_t get_latest_time() const {
+ return queue_.unsync_front()
+ ->column_data(time_col_index_)
+ ->GetValues<int64_t>(1)[latest_ref_row_];
+ }
+
+ bool finished() const { return batches_processed_ == total_batches_; }
+
+ bool advance() {
+ // 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();
+ if (have_active_batch) {
+ // If we have an active batch
+ if (++latest_ref_row_ >= (row_index_t)queue_.unsync_front()->num_rows()) {
+ // hit the end of the batch, need to get the next batch if possible.
+ ++batches_processed_;
+ latest_ref_row_ = 0;
+ have_active_batch &= !queue_.try_pop();
+ if (have_active_batch)
+ assert(queue_.unsync_front()->num_rows() > 0); // empty batches disallowed
+ }
+ }
+ return have_active_batch;
+ }
+
+ // Advance the data to be immediately past the specified TS, updating latest and
+ // latest_ref_row to the latest data prior to that immediate just past Returns true if
+ // updates were made, false if not.
+ bool advance_and_memoize(int64_t ts) {
+ // Advance the right side row index until we reach the latest right row (for each key)
+ // for the given left timestamp.
+
+ // Check if already updated for TS (or if there is no latest)
+ if (empty()) return false; // can't advance if empty
+ auto latest_time = get_latest_time();
+ if (latest_time > ts) return false; // already advanced
+
+ // Not updated. Try to update and possibly advance.
+ bool updated = false;
+ do {
+ latest_time = get_latest_time();
+ // 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 <= ts) {
+ memo_.store(get_latest_batch(), latest_ref_row_, latest_time, get_latest_key());
+ } else {
+ break; // hit a future timestamp -- done updating for now
+ }
+ updated = true;
+ } while (advance());
+ return updated;
+ }
+
+ void push(const std::shared_ptr<arrow::RecordBatch>& rb) {
+ if (rb->num_rows() > 0) {
+ queue_.push(rb);
+ } else {
+ ++batches_processed_; // don't enqueue empty batches, just record as processed
+ }
+ }
+
+ util::optional<const MemoStore::Entry*> get_memo_entry_for_key(KeyType key) {
+ auto r = memo_.get_entry_for_key(key);
+ if (r.has_value()) return r;
+ if (wildcard_key_.has_value()) r = memo_.get_entry_for_key(*wildcard_key_);
+ return r;
+ }
+
+ util::optional<int64_t> get_memo_time_for_key(KeyType key) {
+ auto r = get_memo_entry_for_key(key);
+ return r.has_value() ? util::make_optional((*r)->_time) : util::nullopt;
+ }
+
+ void remove_memo_entries_with_lesser_time(int64_t ts) {
+ memo_.remove_entries_with_lesser_time(ts);
+ }
+
+ const std::shared_ptr<Schema>& get_schema() const { return schema_; }
+
+ void set_total_batches(int n) {
+ assert(n >= 0);
+ assert(total_batches_ == -1); // shouldn't be set more than once
+ total_batches_ = n;
+ }
+
+ private:
+ // Pending record batches. The latest is the front. Batches cannot be empty.
+ ConcurrentQueue<std::shared_ptr<RecordBatch>> queue_;
+
+ // Wildcard key for this input, if applicable.
+ util::optional<KeyType> wildcard_key_;
+
+ // Schema associated with the input
+ std::shared_ptr<Schema> schema_;
+
+ // Total number of batches (only int because InputFinished uses int)
+ int total_batches_ = -1;
+
+ // Number of batches processed so far (only int because InputFinished uses int)
+ int batches_processed_ = 0;
+
+ // Index of the time col
+ col_index_t time_col_index_;
+
+ // Index of the key col
+ col_index_t key_col_index_;
+
+ // Index of the latest row reference within; if >0 then queue_ cannot be empty
+ row_index_t latest_ref_row_ =
+ 0; // must be < queue_.front()->num_rows() if queue_ is non-empty
+
+ // Stores latest known values for the various keys
+ MemoStore memo_;
+
+ // Mapping of source columns to destination columns
+ std::vector<util::optional<col_index_t>> src_to_dst_;
+};
+
+template <size_t MAX_TABLES>
+struct CompositeReferenceRow {
+ struct Entry {
+ arrow::RecordBatch* batch; // can be NULL if there's no value
+ row_index_t row;
+ };
+ Entry refs[MAX_TABLES];
+};
+
+// A table of composite reference rows. Rows maintain pointers to the
+// constituent record batches, but the overall table retains shared_ptr
+// references to ensure memory remains resident while the table is live.
+//
+// The main reason for this is that, especially for wide tables, joins
+// are effectively row-oriented, rather than column-oriented. Separating
+// the join part from the columnar materialization part simplifies the
+// logic around data types and increases efficiency.
+//
+// We don't put the shared_ptr's into the rows for efficiency reasons.
+template <size_t MAX_TABLES>
+class CompositeReferenceTable {
+ public:
+ explicit CompositeReferenceTable(size_t n_tables) : n_tables_(n_tables) {
+ assert(n_tables_ >= 1);
+ assert(n_tables_ <= MAX_TABLES);
+ }
+
+ size_t n_rows() const { return rows_.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<std::unique_ptr<InputState>>& in, int64_t tolerance) {
+ assert(in.size() == n_tables_);
+
+ // Get the LHS key
+ KeyType key = in[0]->get_latest_key();
+
+ // Add row and setup LHS
+ // (the LHS state comes just from the latest row of the LHS table)
+ assert(!in[0]->empty());
+ const std::shared_ptr<arrow::RecordBatch>& lhs_latest_batch =
+ in[0]->get_latest_batch();
+ row_index_t lhs_latest_row = in[0]->get_latest_row();
+ int64_t lhs_latest_time = in[0]->get_latest_time();
+ 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 = rows_.size() + new_batch_size;
+ if (rows_.capacity() < new_capacity) rows_.reserve(new_capacity);
+ }
+ rows_.resize(rows_.size() + 1);
+ auto& row = rows_.back();
+ row.refs[0].batch = lhs_latest_batch.get();
+ row.refs[0].row = lhs_latest_row;
+ add_record_batch_ref(lhs_latest_batch);
+
+ // 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) {
+ util::optional<const MemoStore::Entry*> opt_entry =
+ in[i]->get_memo_entry_for_key(key);
+ if (opt_entry.has_value()) {
+ assert(*opt_entry);
+ if ((*opt_entry)->_time + tolerance >= lhs_latest_time) {
+ // Have a valid entry
+ const MemoStore::Entry* entry = *opt_entry;
+ row.refs[i].batch = entry->_batch.get();
+ row.refs[i].row = entry->_row;
+ add_record_batch_ref(entry->_batch);
+ continue;
+ }
+ }
+ row.refs[i].batch = NULL;
+ row.refs[i].row = 0;
+ }
+ }
+
+ // Materializes the current reference table into a target record batch
+ Result<std::shared_ptr<RecordBatch>> materialize(
Review Comment:
Someone recently pointed me to `src/arrow/table_builder.h` which may be really useful here.
##########
cpp/src/arrow/compute/exec/asof_join_node.cc:
##########
@@ -0,0 +1,800 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+#include <iostream>
+#include <unordered_map>
+
+#include <arrow/api.h>
+#include <arrow/compute/api.h>
+#include <arrow/util/optional.h>
+#include "arrow/compute/exec/asof_join.h"
+#include "arrow/compute/exec/exec_plan.h"
+#include "arrow/compute/exec/options.h"
+#include "arrow/compute/exec/schema_util.h"
+#include "arrow/compute/exec/util.h"
+#include "arrow/result.h"
+#include "arrow/status.h"
+#include "arrow/util/checked_cast.h"
+#include "arrow/util/future.h"
+#include "arrow/util/make_unique.h"
+
+#include <condition_variable>
+#include <mutex>
+#include <thread>
+
+namespace arrow {
+namespace compute {
+
+/**
+ * Simple implementation for an unbound concurrent queue
+ */
+template <class T>
+class ConcurrentQueue {
+ public:
+ T pop() {
+ std::unique_lock<std::mutex> lock(mutex_);
+ cond_.wait(lock, [&] { return !queue_.empty(); });
+ auto item = queue_.front();
+ queue_.pop();
+ return item;
+ }
+
+ void push(const T& item) {
+ std::unique_lock<std::mutex> lock(mutex_);
+ queue_.push(item);
+ cond_.notify_one();
+ }
+
+ util::optional<T> try_pop() {
+ // Try to pop the oldest value from the queue (or return nullopt if none)
+ std::unique_lock<std::mutex> lock(mutex_);
+ if (queue_.empty()) {
+ return util::nullopt;
+ } else {
+ auto item = queue_.front();
+ queue_.pop();
+ return item;
+ }
+ }
+
+ bool empty() const {
+ std::unique_lock<std::mutex> lock(mutex_);
+ return queue_.empty();
+ }
+
+ // Un-synchronized access to front
+ // For this to be "safe":
+ // 1) the caller logically guarantees that queue is not empty
+ // 2) pop/try_pop cannot be called concurrently with this
+ const T& unsync_front() const { return queue_.front(); }
+
+ private:
+ std::queue<T> queue_;
+ mutable std::mutex mutex_;
+ std::condition_variable cond_;
+};
+
+struct MemoStore {
+ // Stores last known values for all the keys
+
+ struct Entry {
+ // Timestamp associated with the entry
+ int64_t _time;
+
+ // Batch associated with the entry (perf is probably OK for this; batches change
+ // rarely)
+ std::shared_ptr<arrow::RecordBatch> _batch;
+
+ // Row associated with the entry
+ row_index_t _row;
+ };
+
+ std::unordered_map<KeyType, Entry> _entries;
+
+ void store(const std::shared_ptr<RecordBatch>& batch, row_index_t row, int64_t time,
+ KeyType key) {
+ auto& e = _entries[key];
+ // that we can do this assignment optionally, is why we
+ // can get array 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;
+ }
+
+ util::optional<const Entry*> get_entry_for_key(KeyType key) const {
+ auto e = _entries.find(key);
+ if (_entries.end() == e) return util::nullopt;
+ return util::optional<const Entry*>(&e->second);
+ }
+
+ void remove_entries_with_lesser_time(int64_t ts) {
+ size_t dbg_size0 = _entries.size();
+ for (auto e = _entries.begin(); e != _entries.end();)
+ if (e->second._time < ts)
+ e = _entries.erase(e);
+ else
+ ++e;
+ size_t dbg_size1 = _entries.size();
+ if (dbg_size1 < dbg_size0) {
+ // cerr << "Removed " << dbg_size0-dbg_size1 << " memo entries.\n";
+ }
+ }
+};
+
+class InputState {
+ // InputState correponds to an input
+ // Input record batches are queued up in InputState until processed and
+ // turned into output record batches.
+
+ public:
+ InputState(const std::shared_ptr<arrow::Schema>& schema,
+ const std::string& time_col_name, const std::string& key_col_name,
+ util::optional<KeyType> wildcard_key)
+ : queue_(),
+ wildcard_key_(wildcard_key),
+ schema_(schema),
+ time_col_index_(
+ schema->GetFieldIndex(time_col_name)), // TODO: handle missing field name
+ key_col_index_(schema->GetFieldIndex(key_col_name)) {}
+
+ col_index_t init_src_to_dst_mapping(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 && is_time_or_key_column(i)))
+ src_to_dst_[i] = dst_offset++;
+ return dst_offset;
+ }
+
+ const util::optional<col_index_t>& map_src_to_dst(col_index_t src) const {
+ return src_to_dst_[src];
+ }
+
+ bool is_time_or_key_column(col_index_t i) const {
+ assert(i < schema_->num_fields());
+ return (i == time_col_index_) || (i == key_col_index_);
+ }
+
+ // Gets the latest row index, assuming the queue isn't empty
+ row_index_t get_latest_row() const { return latest_ref_row_; }
+
+ bool empty() const {
+ if (latest_ref_row_ > 0)
+ return false; // cannot be empty if ref row is >0 -- can avoid slow queue lock
+ // below
+ return queue_.empty();
+ }
+
+ int countbatches_processed_() const { return batches_processed_; }
+ int count_total_batches() const { return total_batches_; }
+
+ // Gets latest batch (precondition: must not be empty)
+ const std::shared_ptr<arrow::RecordBatch>& get_latest_batch() const {
+ return queue_.unsync_front();
+ }
+ KeyType get_latest_key() const {
+ return queue_.unsync_front()
+ ->column_data(key_col_index_)
+ ->GetValues<KeyType>(1)[latest_ref_row_];
+ }
+ int64_t get_latest_time() const {
+ return queue_.unsync_front()
+ ->column_data(time_col_index_)
+ ->GetValues<int64_t>(1)[latest_ref_row_];
+ }
+
+ bool finished() const { return batches_processed_ == total_batches_; }
+
+ bool advance() {
+ // 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();
+ if (have_active_batch) {
+ // If we have an active batch
+ if (++latest_ref_row_ >= (row_index_t)queue_.unsync_front()->num_rows()) {
+ // hit the end of the batch, need to get the next batch if possible.
+ ++batches_processed_;
+ latest_ref_row_ = 0;
+ have_active_batch &= !queue_.try_pop();
+ if (have_active_batch)
+ assert(queue_.unsync_front()->num_rows() > 0); // empty batches disallowed
+ }
+ }
+ return have_active_batch;
+ }
+
+ // Advance the data to be immediately past the specified TS, updating latest and
+ // latest_ref_row to the latest data prior to that immediate just past Returns true if
+ // updates were made, false if not.
+ bool advance_and_memoize(int64_t ts) {
+ // Advance the right side row index until we reach the latest right row (for each key)
+ // for the given left timestamp.
+
+ // Check if already updated for TS (or if there is no latest)
+ if (empty()) return false; // can't advance if empty
+ auto latest_time = get_latest_time();
+ if (latest_time > ts) return false; // already advanced
+
+ // Not updated. Try to update and possibly advance.
+ bool updated = false;
+ do {
+ latest_time = get_latest_time();
+ // 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 <= ts) {
+ memo_.store(get_latest_batch(), latest_ref_row_, latest_time, get_latest_key());
+ } else {
+ break; // hit a future timestamp -- done updating for now
+ }
+ updated = true;
+ } while (advance());
+ return updated;
+ }
+
+ void push(const std::shared_ptr<arrow::RecordBatch>& rb) {
+ if (rb->num_rows() > 0) {
+ queue_.push(rb);
+ } else {
+ ++batches_processed_; // don't enqueue empty batches, just record as processed
+ }
+ }
+
+ util::optional<const MemoStore::Entry*> get_memo_entry_for_key(KeyType key) {
+ auto r = memo_.get_entry_for_key(key);
+ if (r.has_value()) return r;
+ if (wildcard_key_.has_value()) r = memo_.get_entry_for_key(*wildcard_key_);
+ return r;
+ }
+
+ util::optional<int64_t> get_memo_time_for_key(KeyType key) {
+ auto r = get_memo_entry_for_key(key);
+ return r.has_value() ? util::make_optional((*r)->_time) : util::nullopt;
+ }
+
+ void remove_memo_entries_with_lesser_time(int64_t ts) {
+ memo_.remove_entries_with_lesser_time(ts);
+ }
+
+ const std::shared_ptr<Schema>& get_schema() const { return schema_; }
+
+ void set_total_batches(int n) {
+ assert(n >= 0);
+ assert(total_batches_ == -1); // shouldn't be set more than once
+ total_batches_ = n;
+ }
+
+ private:
+ // Pending record batches. The latest is the front. Batches cannot be empty.
+ ConcurrentQueue<std::shared_ptr<RecordBatch>> queue_;
+
+ // Wildcard key for this input, if applicable.
+ util::optional<KeyType> wildcard_key_;
+
+ // Schema associated with the input
+ std::shared_ptr<Schema> schema_;
+
+ // Total number of batches (only int because InputFinished uses int)
+ int total_batches_ = -1;
+
+ // Number of batches processed so far (only int because InputFinished uses int)
+ int batches_processed_ = 0;
+
+ // Index of the time col
+ col_index_t time_col_index_;
+
+ // Index of the key col
+ col_index_t key_col_index_;
+
+ // Index of the latest row reference within; if >0 then queue_ cannot be empty
+ row_index_t latest_ref_row_ =
+ 0; // must be < queue_.front()->num_rows() if queue_ is non-empty
+
+ // Stores latest known values for the various keys
+ MemoStore memo_;
+
+ // Mapping of source columns to destination columns
+ std::vector<util::optional<col_index_t>> src_to_dst_;
+};
+
+template <size_t MAX_TABLES>
+struct CompositeReferenceRow {
+ struct Entry {
+ arrow::RecordBatch* batch; // can be NULL if there's no value
+ row_index_t row;
+ };
+ Entry refs[MAX_TABLES];
+};
+
+// A table of composite reference rows. Rows maintain pointers to the
+// constituent record batches, but the overall table retains shared_ptr
+// references to ensure memory remains resident while the table is live.
+//
+// The main reason for this is that, especially for wide tables, joins
+// are effectively row-oriented, rather than column-oriented. Separating
+// the join part from the columnar materialization part simplifies the
+// logic around data types and increases efficiency.
+//
+// We don't put the shared_ptr's into the rows for efficiency reasons.
+template <size_t MAX_TABLES>
+class CompositeReferenceTable {
+ public:
+ explicit CompositeReferenceTable(size_t n_tables) : n_tables_(n_tables) {
+ assert(n_tables_ >= 1);
+ assert(n_tables_ <= MAX_TABLES);
+ }
+
+ size_t n_rows() const { return rows_.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<std::unique_ptr<InputState>>& in, int64_t tolerance) {
+ assert(in.size() == n_tables_);
+
+ // Get the LHS key
+ KeyType key = in[0]->get_latest_key();
+
+ // Add row and setup LHS
+ // (the LHS state comes just from the latest row of the LHS table)
+ assert(!in[0]->empty());
+ const std::shared_ptr<arrow::RecordBatch>& lhs_latest_batch =
+ in[0]->get_latest_batch();
+ row_index_t lhs_latest_row = in[0]->get_latest_row();
+ int64_t lhs_latest_time = in[0]->get_latest_time();
+ 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 = rows_.size() + new_batch_size;
+ if (rows_.capacity() < new_capacity) rows_.reserve(new_capacity);
+ }
+ rows_.resize(rows_.size() + 1);
+ auto& row = rows_.back();
+ row.refs[0].batch = lhs_latest_batch.get();
+ row.refs[0].row = lhs_latest_row;
+ add_record_batch_ref(lhs_latest_batch);
+
+ // 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) {
+ util::optional<const MemoStore::Entry*> opt_entry =
+ in[i]->get_memo_entry_for_key(key);
+ if (opt_entry.has_value()) {
+ assert(*opt_entry);
+ if ((*opt_entry)->_time + tolerance >= lhs_latest_time) {
+ // Have a valid entry
+ const MemoStore::Entry* entry = *opt_entry;
+ row.refs[i].batch = entry->_batch.get();
+ row.refs[i].row = entry->_row;
+ add_record_batch_ref(entry->_batch);
+ continue;
+ }
+ }
+ row.refs[i].batch = NULL;
+ row.refs[i].row = 0;
+ }
+ }
+
+ // Materializes the current reference table into a target record batch
+ Result<std::shared_ptr<RecordBatch>> materialize(
+ const std::shared_ptr<arrow::Schema>& output_schema,
+ const std::vector<std::unique_ptr<InputState>>& state) {
+ // cerr << "materialize BEGIN\n";
+ assert(state.size() == n_tables_);
+ assert(state.size() >= 1);
+
+ // Don't build empty batches
+ size_t n_rows = rows_.size();
+ if (!n_rows) return NULLPTR;
+
+ // Build the arrays column-by-column from the rows
+ std::vector<std::shared_ptr<arrow::Array>> arrays(output_schema->num_fields());
+ for (size_t i_table = 0; i_table < n_tables_; ++i_table) {
+ int n_src_cols = state.at(i_table)->get_schema()->num_fields();
+ {
+ for (col_index_t i_src_col = 0; i_src_col < n_src_cols; ++i_src_col) {
+ util::optional<col_index_t> i_dst_col_opt =
+ state[i_table]->map_src_to_dst(i_src_col);
+ if (!i_dst_col_opt) continue;
+ col_index_t i_dst_col = *i_dst_col_opt;
+ const auto& src_field = state[i_table]->get_schema()->field(i_src_col);
+ const auto& dst_field = output_schema->field(i_dst_col);
+ assert(src_field->type()->Equals(dst_field->type()));
+ assert(src_field->name() == dst_field->name());
+ const auto& field_type = src_field->type();
+
+ if (field_type->Equals(arrow::int32())) {
+ ARROW_ASSIGN_OR_RAISE(
+ arrays.at(i_dst_col),
+ (materialize_primitive_column<arrow::Int32Builder, int32_t>(i_table,
+ i_src_col)));
+ } else if (field_type->Equals(arrow::int64())) {
+ ARROW_ASSIGN_OR_RAISE(
+ arrays.at(i_dst_col),
+ (materialize_primitive_column<arrow::Int64Builder, int64_t>(i_table,
+ i_src_col)));
+ } else if (field_type->Equals(arrow::float64())) {
+ ARROW_ASSIGN_OR_RAISE(
+ arrays.at(i_dst_col),
+ (materialize_primitive_column<arrow::DoubleBuilder, double>(i_table,
+ i_src_col)));
+ } else {
+ ARROW_RETURN_NOT_OK(
+ Status::Invalid("Unsupported data type: ", src_field->name()));
+ }
+ }
+ }
+ }
+
+ // Build the result
+ assert(sizeof(size_t) >= sizeof(int64_t)); // Make takes signed int64_t for num_rows
+
+ // TODO: check n_rows for cast
Review Comment:
Create a follow-up JIRA and then change this comment to `TODO(ARROW-XYZ): `
##########
cpp/src/arrow/compute/exec/asof_join_node.cc:
##########
@@ -0,0 +1,800 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+#include <iostream>
+#include <unordered_map>
+
+#include <arrow/api.h>
+#include <arrow/compute/api.h>
+#include <arrow/util/optional.h>
+#include "arrow/compute/exec/asof_join.h"
+#include "arrow/compute/exec/exec_plan.h"
+#include "arrow/compute/exec/options.h"
+#include "arrow/compute/exec/schema_util.h"
+#include "arrow/compute/exec/util.h"
+#include "arrow/result.h"
+#include "arrow/status.h"
+#include "arrow/util/checked_cast.h"
+#include "arrow/util/future.h"
+#include "arrow/util/make_unique.h"
+
+#include <condition_variable>
+#include <mutex>
+#include <thread>
+
+namespace arrow {
+namespace compute {
+
+/**
+ * Simple implementation for an unbound concurrent queue
+ */
+template <class T>
+class ConcurrentQueue {
+ public:
+ T pop() {
+ std::unique_lock<std::mutex> lock(mutex_);
+ cond_.wait(lock, [&] { return !queue_.empty(); });
+ auto item = queue_.front();
+ queue_.pop();
+ return item;
+ }
+
+ void push(const T& item) {
+ std::unique_lock<std::mutex> lock(mutex_);
+ queue_.push(item);
+ cond_.notify_one();
+ }
+
+ util::optional<T> try_pop() {
+ // Try to pop the oldest value from the queue (or return nullopt if none)
+ std::unique_lock<std::mutex> lock(mutex_);
+ if (queue_.empty()) {
+ return util::nullopt;
+ } else {
+ auto item = queue_.front();
+ queue_.pop();
+ return item;
+ }
+ }
+
+ bool empty() const {
+ std::unique_lock<std::mutex> lock(mutex_);
+ return queue_.empty();
+ }
+
+ // Un-synchronized access to front
+ // For this to be "safe":
+ // 1) the caller logically guarantees that queue is not empty
+ // 2) pop/try_pop cannot be called concurrently with this
+ const T& unsync_front() const { return queue_.front(); }
+
+ private:
+ std::queue<T> queue_;
+ mutable std::mutex mutex_;
+ std::condition_variable cond_;
+};
+
+struct MemoStore {
+ // Stores last known values for all the keys
+
+ struct Entry {
+ // Timestamp associated with the entry
+ int64_t _time;
+
+ // Batch associated with the entry (perf is probably OK for this; batches change
+ // rarely)
+ std::shared_ptr<arrow::RecordBatch> _batch;
+
+ // Row associated with the entry
+ row_index_t _row;
+ };
+
+ std::unordered_map<KeyType, Entry> _entries;
+
+ void store(const std::shared_ptr<RecordBatch>& batch, row_index_t row, int64_t time,
+ KeyType key) {
+ auto& e = _entries[key];
+ // that we can do this assignment optionally, is why we
+ // can get array 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;
+ }
+
+ util::optional<const Entry*> get_entry_for_key(KeyType key) const {
+ auto e = _entries.find(key);
+ if (_entries.end() == e) return util::nullopt;
+ return util::optional<const Entry*>(&e->second);
+ }
+
+ void remove_entries_with_lesser_time(int64_t ts) {
+ size_t dbg_size0 = _entries.size();
+ for (auto e = _entries.begin(); e != _entries.end();)
+ if (e->second._time < ts)
+ e = _entries.erase(e);
+ else
+ ++e;
+ size_t dbg_size1 = _entries.size();
+ if (dbg_size1 < dbg_size0) {
+ // cerr << "Removed " << dbg_size0-dbg_size1 << " memo entries.\n";
+ }
+ }
+};
+
+class InputState {
+ // InputState correponds to an input
+ // Input record batches are queued up in InputState until processed and
+ // turned into output record batches.
+
+ public:
+ InputState(const std::shared_ptr<arrow::Schema>& schema,
+ const std::string& time_col_name, const std::string& key_col_name,
+ util::optional<KeyType> wildcard_key)
+ : queue_(),
+ wildcard_key_(wildcard_key),
+ schema_(schema),
+ time_col_index_(
+ schema->GetFieldIndex(time_col_name)), // TODO: handle missing field name
+ key_col_index_(schema->GetFieldIndex(key_col_name)) {}
+
+ col_index_t init_src_to_dst_mapping(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 && is_time_or_key_column(i)))
+ src_to_dst_[i] = dst_offset++;
+ return dst_offset;
+ }
+
+ const util::optional<col_index_t>& map_src_to_dst(col_index_t src) const {
+ return src_to_dst_[src];
+ }
+
+ bool is_time_or_key_column(col_index_t i) const {
+ assert(i < schema_->num_fields());
+ return (i == time_col_index_) || (i == key_col_index_);
+ }
+
+ // Gets the latest row index, assuming the queue isn't empty
+ row_index_t get_latest_row() const { return latest_ref_row_; }
+
+ bool empty() const {
+ if (latest_ref_row_ > 0)
+ return false; // cannot be empty if ref row is >0 -- can avoid slow queue lock
+ // below
+ return queue_.empty();
+ }
+
+ int countbatches_processed_() const { return batches_processed_; }
+ int count_total_batches() const { return total_batches_; }
+
+ // Gets latest batch (precondition: must not be empty)
+ const std::shared_ptr<arrow::RecordBatch>& get_latest_batch() const {
+ return queue_.unsync_front();
+ }
+ KeyType get_latest_key() const {
+ return queue_.unsync_front()
+ ->column_data(key_col_index_)
+ ->GetValues<KeyType>(1)[latest_ref_row_];
+ }
+ int64_t get_latest_time() const {
+ return queue_.unsync_front()
+ ->column_data(time_col_index_)
+ ->GetValues<int64_t>(1)[latest_ref_row_];
+ }
+
+ bool finished() const { return batches_processed_ == total_batches_; }
+
+ bool advance() {
+ // 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();
+ if (have_active_batch) {
+ // If we have an active batch
+ if (++latest_ref_row_ >= (row_index_t)queue_.unsync_front()->num_rows()) {
+ // hit the end of the batch, need to get the next batch if possible.
+ ++batches_processed_;
+ latest_ref_row_ = 0;
+ have_active_batch &= !queue_.try_pop();
+ if (have_active_batch)
+ assert(queue_.unsync_front()->num_rows() > 0); // empty batches disallowed
+ }
+ }
+ return have_active_batch;
+ }
+
+ // Advance the data to be immediately past the specified TS, updating latest and
+ // latest_ref_row to the latest data prior to that immediate just past Returns true if
+ // updates were made, false if not.
+ bool advance_and_memoize(int64_t ts) {
+ // Advance the right side row index until we reach the latest right row (for each key)
+ // for the given left timestamp.
+
+ // Check if already updated for TS (or if there is no latest)
+ if (empty()) return false; // can't advance if empty
+ auto latest_time = get_latest_time();
+ if (latest_time > ts) return false; // already advanced
+
+ // Not updated. Try to update and possibly advance.
+ bool updated = false;
+ do {
+ latest_time = get_latest_time();
+ // 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 <= ts) {
+ memo_.store(get_latest_batch(), latest_ref_row_, latest_time, get_latest_key());
+ } else {
+ break; // hit a future timestamp -- done updating for now
+ }
+ updated = true;
+ } while (advance());
+ return updated;
+ }
+
+ void push(const std::shared_ptr<arrow::RecordBatch>& rb) {
+ if (rb->num_rows() > 0) {
+ queue_.push(rb);
+ } else {
+ ++batches_processed_; // don't enqueue empty batches, just record as processed
+ }
+ }
+
+ util::optional<const MemoStore::Entry*> get_memo_entry_for_key(KeyType key) {
+ auto r = memo_.get_entry_for_key(key);
+ if (r.has_value()) return r;
+ if (wildcard_key_.has_value()) r = memo_.get_entry_for_key(*wildcard_key_);
+ return r;
+ }
+
+ util::optional<int64_t> get_memo_time_for_key(KeyType key) {
+ auto r = get_memo_entry_for_key(key);
+ return r.has_value() ? util::make_optional((*r)->_time) : util::nullopt;
+ }
+
+ void remove_memo_entries_with_lesser_time(int64_t ts) {
+ memo_.remove_entries_with_lesser_time(ts);
+ }
+
+ const std::shared_ptr<Schema>& get_schema() const { return schema_; }
+
+ void set_total_batches(int n) {
+ assert(n >= 0);
+ assert(total_batches_ == -1); // shouldn't be set more than once
+ total_batches_ = n;
+ }
+
+ private:
+ // Pending record batches. The latest is the front. Batches cannot be empty.
+ ConcurrentQueue<std::shared_ptr<RecordBatch>> queue_;
+
+ // Wildcard key for this input, if applicable.
+ util::optional<KeyType> wildcard_key_;
+
+ // Schema associated with the input
+ std::shared_ptr<Schema> schema_;
+
+ // Total number of batches (only int because InputFinished uses int)
+ int total_batches_ = -1;
+
+ // Number of batches processed so far (only int because InputFinished uses int)
+ int batches_processed_ = 0;
+
+ // Index of the time col
+ col_index_t time_col_index_;
+
+ // Index of the key col
+ col_index_t key_col_index_;
+
+ // Index of the latest row reference within; if >0 then queue_ cannot be empty
+ row_index_t latest_ref_row_ =
+ 0; // must be < queue_.front()->num_rows() if queue_ is non-empty
+
+ // Stores latest known values for the various keys
+ MemoStore memo_;
+
+ // Mapping of source columns to destination columns
+ std::vector<util::optional<col_index_t>> src_to_dst_;
+};
+
+template <size_t MAX_TABLES>
+struct CompositeReferenceRow {
+ struct Entry {
+ arrow::RecordBatch* batch; // can be NULL if there's no value
+ row_index_t row;
+ };
+ Entry refs[MAX_TABLES];
+};
+
+// A table of composite reference rows. Rows maintain pointers to the
+// constituent record batches, but the overall table retains shared_ptr
+// references to ensure memory remains resident while the table is live.
+//
+// The main reason for this is that, especially for wide tables, joins
+// are effectively row-oriented, rather than column-oriented. Separating
+// the join part from the columnar materialization part simplifies the
+// logic around data types and increases efficiency.
+//
+// We don't put the shared_ptr's into the rows for efficiency reasons.
+template <size_t MAX_TABLES>
+class CompositeReferenceTable {
+ public:
+ explicit CompositeReferenceTable(size_t n_tables) : n_tables_(n_tables) {
+ assert(n_tables_ >= 1);
+ assert(n_tables_ <= MAX_TABLES);
+ }
+
+ size_t n_rows() const { return rows_.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<std::unique_ptr<InputState>>& in, int64_t tolerance) {
+ assert(in.size() == n_tables_);
+
+ // Get the LHS key
+ KeyType key = in[0]->get_latest_key();
+
+ // Add row and setup LHS
+ // (the LHS state comes just from the latest row of the LHS table)
+ assert(!in[0]->empty());
+ const std::shared_ptr<arrow::RecordBatch>& lhs_latest_batch =
+ in[0]->get_latest_batch();
+ row_index_t lhs_latest_row = in[0]->get_latest_row();
+ int64_t lhs_latest_time = in[0]->get_latest_time();
+ 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 = rows_.size() + new_batch_size;
+ if (rows_.capacity() < new_capacity) rows_.reserve(new_capacity);
+ }
+ rows_.resize(rows_.size() + 1);
+ auto& row = rows_.back();
+ row.refs[0].batch = lhs_latest_batch.get();
+ row.refs[0].row = lhs_latest_row;
+ add_record_batch_ref(lhs_latest_batch);
+
+ // 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) {
+ util::optional<const MemoStore::Entry*> opt_entry =
+ in[i]->get_memo_entry_for_key(key);
+ if (opt_entry.has_value()) {
+ assert(*opt_entry);
+ if ((*opt_entry)->_time + tolerance >= lhs_latest_time) {
+ // Have a valid entry
+ const MemoStore::Entry* entry = *opt_entry;
+ row.refs[i].batch = entry->_batch.get();
+ row.refs[i].row = entry->_row;
+ add_record_batch_ref(entry->_batch);
+ continue;
+ }
+ }
+ row.refs[i].batch = NULL;
+ row.refs[i].row = 0;
+ }
+ }
+
+ // Materializes the current reference table into a target record batch
+ Result<std::shared_ptr<RecordBatch>> materialize(
+ const std::shared_ptr<arrow::Schema>& output_schema,
+ const std::vector<std::unique_ptr<InputState>>& state) {
+ // cerr << "materialize BEGIN\n";
+ assert(state.size() == n_tables_);
+ assert(state.size() >= 1);
+
+ // Don't build empty batches
+ size_t n_rows = rows_.size();
+ if (!n_rows) return NULLPTR;
+
+ // Build the arrays column-by-column from the rows
+ std::vector<std::shared_ptr<arrow::Array>> arrays(output_schema->num_fields());
+ for (size_t i_table = 0; i_table < n_tables_; ++i_table) {
+ int n_src_cols = state.at(i_table)->get_schema()->num_fields();
+ {
+ for (col_index_t i_src_col = 0; i_src_col < n_src_cols; ++i_src_col) {
+ util::optional<col_index_t> i_dst_col_opt =
+ state[i_table]->map_src_to_dst(i_src_col);
+ if (!i_dst_col_opt) continue;
+ col_index_t i_dst_col = *i_dst_col_opt;
+ const auto& src_field = state[i_table]->get_schema()->field(i_src_col);
+ const auto& dst_field = output_schema->field(i_dst_col);
+ assert(src_field->type()->Equals(dst_field->type()));
+ assert(src_field->name() == dst_field->name());
Review Comment:
```suggestion
DCHECK(src_field->type()->Equals(dst_field->type()));
DCHECK_EQ(src_field->name(), dst_field->name());
```
##########
cpp/src/arrow/compute/exec/asof_join_node.cc:
##########
@@ -0,0 +1,800 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+#include <iostream>
+#include <unordered_map>
+
+#include <arrow/api.h>
+#include <arrow/compute/api.h>
+#include <arrow/util/optional.h>
+#include "arrow/compute/exec/asof_join.h"
+#include "arrow/compute/exec/exec_plan.h"
+#include "arrow/compute/exec/options.h"
+#include "arrow/compute/exec/schema_util.h"
+#include "arrow/compute/exec/util.h"
+#include "arrow/result.h"
+#include "arrow/status.h"
+#include "arrow/util/checked_cast.h"
+#include "arrow/util/future.h"
+#include "arrow/util/make_unique.h"
+
+#include <condition_variable>
+#include <mutex>
+#include <thread>
+
+namespace arrow {
+namespace compute {
+
+/**
+ * Simple implementation for an unbound concurrent queue
+ */
+template <class T>
+class ConcurrentQueue {
+ public:
+ T pop() {
+ std::unique_lock<std::mutex> lock(mutex_);
+ cond_.wait(lock, [&] { return !queue_.empty(); });
+ auto item = queue_.front();
+ queue_.pop();
+ return item;
+ }
+
+ void push(const T& item) {
+ std::unique_lock<std::mutex> lock(mutex_);
+ queue_.push(item);
+ cond_.notify_one();
+ }
+
+ util::optional<T> try_pop() {
+ // Try to pop the oldest value from the queue (or return nullopt if none)
+ std::unique_lock<std::mutex> lock(mutex_);
+ if (queue_.empty()) {
+ return util::nullopt;
+ } else {
+ auto item = queue_.front();
+ queue_.pop();
+ return item;
+ }
+ }
+
+ bool empty() const {
+ std::unique_lock<std::mutex> lock(mutex_);
+ return queue_.empty();
+ }
+
+ // Un-synchronized access to front
+ // For this to be "safe":
+ // 1) the caller logically guarantees that queue is not empty
+ // 2) pop/try_pop cannot be called concurrently with this
+ const T& unsync_front() const { return queue_.front(); }
+
+ private:
+ std::queue<T> queue_;
+ mutable std::mutex mutex_;
+ std::condition_variable cond_;
+};
+
+struct MemoStore {
+ // Stores last known values for all the keys
+
+ struct Entry {
+ // Timestamp associated with the entry
+ int64_t _time;
+
+ // Batch associated with the entry (perf is probably OK for this; batches change
+ // rarely)
+ std::shared_ptr<arrow::RecordBatch> _batch;
+
+ // Row associated with the entry
+ row_index_t _row;
+ };
+
+ std::unordered_map<KeyType, Entry> _entries;
+
+ void store(const std::shared_ptr<RecordBatch>& batch, row_index_t row, int64_t time,
+ KeyType key) {
+ auto& e = _entries[key];
+ // that we can do this assignment optionally, is why we
+ // can get array 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;
+ }
+
+ util::optional<const Entry*> get_entry_for_key(KeyType key) const {
+ auto e = _entries.find(key);
+ if (_entries.end() == e) return util::nullopt;
+ return util::optional<const Entry*>(&e->second);
+ }
+
+ void remove_entries_with_lesser_time(int64_t ts) {
+ size_t dbg_size0 = _entries.size();
+ for (auto e = _entries.begin(); e != _entries.end();)
+ if (e->second._time < ts)
+ e = _entries.erase(e);
+ else
+ ++e;
+ size_t dbg_size1 = _entries.size();
+ if (dbg_size1 < dbg_size0) {
+ // cerr << "Removed " << dbg_size0-dbg_size1 << " memo entries.\n";
+ }
+ }
+};
+
+class InputState {
+ // InputState correponds to an input
+ // Input record batches are queued up in InputState until processed and
+ // turned into output record batches.
+
+ public:
+ InputState(const std::shared_ptr<arrow::Schema>& schema,
+ const std::string& time_col_name, const std::string& key_col_name,
+ util::optional<KeyType> wildcard_key)
+ : queue_(),
+ wildcard_key_(wildcard_key),
+ schema_(schema),
+ time_col_index_(
+ schema->GetFieldIndex(time_col_name)), // TODO: handle missing field name
+ key_col_index_(schema->GetFieldIndex(key_col_name)) {}
+
+ col_index_t init_src_to_dst_mapping(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 && is_time_or_key_column(i)))
+ src_to_dst_[i] = dst_offset++;
+ return dst_offset;
+ }
+
+ const util::optional<col_index_t>& map_src_to_dst(col_index_t src) const {
+ return src_to_dst_[src];
+ }
+
+ bool is_time_or_key_column(col_index_t i) const {
+ assert(i < schema_->num_fields());
+ return (i == time_col_index_) || (i == key_col_index_);
+ }
+
+ // Gets the latest row index, assuming the queue isn't empty
+ row_index_t get_latest_row() const { return latest_ref_row_; }
+
+ bool empty() const {
+ if (latest_ref_row_ > 0)
+ return false; // cannot be empty if ref row is >0 -- can avoid slow queue lock
+ // below
+ return queue_.empty();
+ }
+
+ int countbatches_processed_() const { return batches_processed_; }
+ int count_total_batches() const { return total_batches_; }
+
+ // Gets latest batch (precondition: must not be empty)
+ const std::shared_ptr<arrow::RecordBatch>& get_latest_batch() const {
+ return queue_.unsync_front();
+ }
+ KeyType get_latest_key() const {
+ return queue_.unsync_front()
+ ->column_data(key_col_index_)
+ ->GetValues<KeyType>(1)[latest_ref_row_];
+ }
+ int64_t get_latest_time() const {
+ return queue_.unsync_front()
+ ->column_data(time_col_index_)
+ ->GetValues<int64_t>(1)[latest_ref_row_];
+ }
+
+ bool finished() const { return batches_processed_ == total_batches_; }
+
+ bool advance() {
+ // 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();
+ if (have_active_batch) {
+ // If we have an active batch
+ if (++latest_ref_row_ >= (row_index_t)queue_.unsync_front()->num_rows()) {
+ // hit the end of the batch, need to get the next batch if possible.
+ ++batches_processed_;
+ latest_ref_row_ = 0;
+ have_active_batch &= !queue_.try_pop();
+ if (have_active_batch)
+ assert(queue_.unsync_front()->num_rows() > 0); // empty batches disallowed
+ }
+ }
+ return have_active_batch;
+ }
+
+ // Advance the data to be immediately past the specified TS, updating latest and
+ // latest_ref_row to the latest data prior to that immediate just past Returns true if
+ // updates were made, false if not.
+ bool advance_and_memoize(int64_t ts) {
+ // Advance the right side row index until we reach the latest right row (for each key)
+ // for the given left timestamp.
+
+ // Check if already updated for TS (or if there is no latest)
+ if (empty()) return false; // can't advance if empty
+ auto latest_time = get_latest_time();
+ if (latest_time > ts) return false; // already advanced
+
+ // Not updated. Try to update and possibly advance.
+ bool updated = false;
+ do {
+ latest_time = get_latest_time();
+ // 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 <= ts) {
+ memo_.store(get_latest_batch(), latest_ref_row_, latest_time, get_latest_key());
+ } else {
+ break; // hit a future timestamp -- done updating for now
+ }
+ updated = true;
+ } while (advance());
+ return updated;
+ }
+
+ void push(const std::shared_ptr<arrow::RecordBatch>& rb) {
+ if (rb->num_rows() > 0) {
+ queue_.push(rb);
+ } else {
+ ++batches_processed_; // don't enqueue empty batches, just record as processed
+ }
+ }
+
+ util::optional<const MemoStore::Entry*> get_memo_entry_for_key(KeyType key) {
+ auto r = memo_.get_entry_for_key(key);
+ if (r.has_value()) return r;
+ if (wildcard_key_.has_value()) r = memo_.get_entry_for_key(*wildcard_key_);
+ return r;
+ }
+
+ util::optional<int64_t> get_memo_time_for_key(KeyType key) {
+ auto r = get_memo_entry_for_key(key);
+ return r.has_value() ? util::make_optional((*r)->_time) : util::nullopt;
+ }
+
+ void remove_memo_entries_with_lesser_time(int64_t ts) {
+ memo_.remove_entries_with_lesser_time(ts);
+ }
+
+ const std::shared_ptr<Schema>& get_schema() const { return schema_; }
+
+ void set_total_batches(int n) {
+ assert(n >= 0);
+ assert(total_batches_ == -1); // shouldn't be set more than once
+ total_batches_ = n;
+ }
+
+ private:
+ // Pending record batches. The latest is the front. Batches cannot be empty.
+ ConcurrentQueue<std::shared_ptr<RecordBatch>> queue_;
+
+ // Wildcard key for this input, if applicable.
+ util::optional<KeyType> wildcard_key_;
+
+ // Schema associated with the input
+ std::shared_ptr<Schema> schema_;
+
+ // Total number of batches (only int because InputFinished uses int)
+ int total_batches_ = -1;
+
+ // Number of batches processed so far (only int because InputFinished uses int)
+ int batches_processed_ = 0;
+
+ // Index of the time col
+ col_index_t time_col_index_;
+
+ // Index of the key col
+ col_index_t key_col_index_;
+
+ // Index of the latest row reference within; if >0 then queue_ cannot be empty
+ row_index_t latest_ref_row_ =
+ 0; // must be < queue_.front()->num_rows() if queue_ is non-empty
+
+ // Stores latest known values for the various keys
+ MemoStore memo_;
+
+ // Mapping of source columns to destination columns
+ std::vector<util::optional<col_index_t>> src_to_dst_;
+};
+
+template <size_t MAX_TABLES>
+struct CompositeReferenceRow {
+ struct Entry {
+ arrow::RecordBatch* batch; // can be NULL if there's no value
+ row_index_t row;
+ };
+ Entry refs[MAX_TABLES];
+};
+
+// A table of composite reference rows. Rows maintain pointers to the
+// constituent record batches, but the overall table retains shared_ptr
+// references to ensure memory remains resident while the table is live.
+//
+// The main reason for this is that, especially for wide tables, joins
+// are effectively row-oriented, rather than column-oriented. Separating
+// the join part from the columnar materialization part simplifies the
+// logic around data types and increases efficiency.
+//
+// We don't put the shared_ptr's into the rows for efficiency reasons.
+template <size_t MAX_TABLES>
+class CompositeReferenceTable {
+ public:
+ explicit CompositeReferenceTable(size_t n_tables) : n_tables_(n_tables) {
+ assert(n_tables_ >= 1);
+ assert(n_tables_ <= MAX_TABLES);
+ }
+
+ size_t n_rows() const { return rows_.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<std::unique_ptr<InputState>>& in, int64_t tolerance) {
+ assert(in.size() == n_tables_);
+
+ // Get the LHS key
+ KeyType key = in[0]->get_latest_key();
+
+ // Add row and setup LHS
+ // (the LHS state comes just from the latest row of the LHS table)
+ assert(!in[0]->empty());
+ const std::shared_ptr<arrow::RecordBatch>& lhs_latest_batch =
+ in[0]->get_latest_batch();
+ row_index_t lhs_latest_row = in[0]->get_latest_row();
+ int64_t lhs_latest_time = in[0]->get_latest_time();
+ 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 = rows_.size() + new_batch_size;
+ if (rows_.capacity() < new_capacity) rows_.reserve(new_capacity);
+ }
+ rows_.resize(rows_.size() + 1);
+ auto& row = rows_.back();
+ row.refs[0].batch = lhs_latest_batch.get();
+ row.refs[0].row = lhs_latest_row;
+ add_record_batch_ref(lhs_latest_batch);
+
+ // 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) {
+ util::optional<const MemoStore::Entry*> opt_entry =
+ in[i]->get_memo_entry_for_key(key);
+ if (opt_entry.has_value()) {
+ assert(*opt_entry);
+ if ((*opt_entry)->_time + tolerance >= lhs_latest_time) {
+ // Have a valid entry
+ const MemoStore::Entry* entry = *opt_entry;
+ row.refs[i].batch = entry->_batch.get();
+ row.refs[i].row = entry->_row;
+ add_record_batch_ref(entry->_batch);
+ continue;
+ }
+ }
+ row.refs[i].batch = NULL;
+ row.refs[i].row = 0;
+ }
+ }
+
+ // Materializes the current reference table into a target record batch
+ Result<std::shared_ptr<RecordBatch>> materialize(
+ const std::shared_ptr<arrow::Schema>& output_schema,
+ const std::vector<std::unique_ptr<InputState>>& state) {
+ // cerr << "materialize BEGIN\n";
+ assert(state.size() == n_tables_);
+ assert(state.size() >= 1);
+
+ // Don't build empty batches
+ size_t n_rows = rows_.size();
+ if (!n_rows) return NULLPTR;
+
+ // Build the arrays column-by-column from the rows
+ std::vector<std::shared_ptr<arrow::Array>> arrays(output_schema->num_fields());
+ for (size_t i_table = 0; i_table < n_tables_; ++i_table) {
+ int n_src_cols = state.at(i_table)->get_schema()->num_fields();
+ {
+ for (col_index_t i_src_col = 0; i_src_col < n_src_cols; ++i_src_col) {
+ util::optional<col_index_t> i_dst_col_opt =
+ state[i_table]->map_src_to_dst(i_src_col);
+ if (!i_dst_col_opt) continue;
+ col_index_t i_dst_col = *i_dst_col_opt;
+ const auto& src_field = state[i_table]->get_schema()->field(i_src_col);
+ const auto& dst_field = output_schema->field(i_dst_col);
+ assert(src_field->type()->Equals(dst_field->type()));
+ assert(src_field->name() == dst_field->name());
+ const auto& field_type = src_field->type();
+
+ if (field_type->Equals(arrow::int32())) {
+ ARROW_ASSIGN_OR_RAISE(
+ arrays.at(i_dst_col),
+ (materialize_primitive_column<arrow::Int32Builder, int32_t>(i_table,
+ i_src_col)));
+ } else if (field_type->Equals(arrow::int64())) {
+ ARROW_ASSIGN_OR_RAISE(
+ arrays.at(i_dst_col),
+ (materialize_primitive_column<arrow::Int64Builder, int64_t>(i_table,
+ i_src_col)));
+ } else if (field_type->Equals(arrow::float64())) {
+ ARROW_ASSIGN_OR_RAISE(
+ arrays.at(i_dst_col),
+ (materialize_primitive_column<arrow::DoubleBuilder, double>(i_table,
+ i_src_col)));
+ } else {
+ ARROW_RETURN_NOT_OK(
+ Status::Invalid("Unsupported data type: ", src_field->name()));
+ }
+ }
+ }
+ }
+
+ // Build the result
+ assert(sizeof(size_t) >= sizeof(int64_t)); // Make takes signed int64_t for num_rows
Review Comment:
```suggestion
// Build the result
DCHECK_GE(sizeof(size_t), sizeof(int64_t)) << "AsofJoinNode requires size_t >= 8 bytes";
```
##########
cpp/src/arrow/compute/exec/asof_join_node.cc:
##########
@@ -0,0 +1,800 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+#include <iostream>
+#include <unordered_map>
+
+#include <arrow/api.h>
+#include <arrow/compute/api.h>
+#include <arrow/util/optional.h>
+#include "arrow/compute/exec/asof_join.h"
+#include "arrow/compute/exec/exec_plan.h"
+#include "arrow/compute/exec/options.h"
+#include "arrow/compute/exec/schema_util.h"
+#include "arrow/compute/exec/util.h"
+#include "arrow/result.h"
+#include "arrow/status.h"
+#include "arrow/util/checked_cast.h"
+#include "arrow/util/future.h"
+#include "arrow/util/make_unique.h"
+
+#include <condition_variable>
+#include <mutex>
+#include <thread>
+
+namespace arrow {
+namespace compute {
+
+/**
+ * Simple implementation for an unbound concurrent queue
+ */
+template <class T>
+class ConcurrentQueue {
+ public:
+ T pop() {
+ std::unique_lock<std::mutex> lock(mutex_);
+ cond_.wait(lock, [&] { return !queue_.empty(); });
+ auto item = queue_.front();
+ queue_.pop();
+ return item;
+ }
+
+ void push(const T& item) {
+ std::unique_lock<std::mutex> lock(mutex_);
+ queue_.push(item);
+ cond_.notify_one();
+ }
+
+ util::optional<T> try_pop() {
+ // Try to pop the oldest value from the queue (or return nullopt if none)
+ std::unique_lock<std::mutex> lock(mutex_);
+ if (queue_.empty()) {
+ return util::nullopt;
+ } else {
+ auto item = queue_.front();
+ queue_.pop();
+ return item;
+ }
+ }
+
+ bool empty() const {
+ std::unique_lock<std::mutex> lock(mutex_);
+ return queue_.empty();
+ }
+
+ // Un-synchronized access to front
+ // For this to be "safe":
+ // 1) the caller logically guarantees that queue is not empty
+ // 2) pop/try_pop cannot be called concurrently with this
+ const T& unsync_front() const { return queue_.front(); }
+
+ private:
+ std::queue<T> queue_;
+ mutable std::mutex mutex_;
+ std::condition_variable cond_;
+};
+
+struct MemoStore {
+ // Stores last known values for all the keys
+
+ struct Entry {
+ // Timestamp associated with the entry
+ int64_t _time;
+
+ // Batch associated with the entry (perf is probably OK for this; batches change
+ // rarely)
+ std::shared_ptr<arrow::RecordBatch> _batch;
+
+ // Row associated with the entry
+ row_index_t _row;
+ };
+
+ std::unordered_map<KeyType, Entry> _entries;
+
+ void store(const std::shared_ptr<RecordBatch>& batch, row_index_t row, int64_t time,
+ KeyType key) {
+ auto& e = _entries[key];
+ // that we can do this assignment optionally, is why we
+ // can get array 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;
+ }
+
+ util::optional<const Entry*> get_entry_for_key(KeyType key) const {
+ auto e = _entries.find(key);
+ if (_entries.end() == e) return util::nullopt;
+ return util::optional<const Entry*>(&e->second);
+ }
+
+ void remove_entries_with_lesser_time(int64_t ts) {
+ size_t dbg_size0 = _entries.size();
+ for (auto e = _entries.begin(); e != _entries.end();)
+ if (e->second._time < ts)
+ e = _entries.erase(e);
+ else
+ ++e;
+ size_t dbg_size1 = _entries.size();
+ if (dbg_size1 < dbg_size0) {
+ // cerr << "Removed " << dbg_size0-dbg_size1 << " memo entries.\n";
+ }
+ }
+};
+
+class InputState {
+ // InputState correponds to an input
+ // Input record batches are queued up in InputState until processed and
+ // turned into output record batches.
+
+ public:
+ InputState(const std::shared_ptr<arrow::Schema>& schema,
+ const std::string& time_col_name, const std::string& key_col_name,
+ util::optional<KeyType> wildcard_key)
+ : queue_(),
+ wildcard_key_(wildcard_key),
+ schema_(schema),
+ time_col_index_(
+ schema->GetFieldIndex(time_col_name)), // TODO: handle missing field name
+ key_col_index_(schema->GetFieldIndex(key_col_name)) {}
+
+ col_index_t init_src_to_dst_mapping(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 && is_time_or_key_column(i)))
+ src_to_dst_[i] = dst_offset++;
+ return dst_offset;
+ }
+
+ const util::optional<col_index_t>& map_src_to_dst(col_index_t src) const {
+ return src_to_dst_[src];
+ }
+
+ bool is_time_or_key_column(col_index_t i) const {
+ assert(i < schema_->num_fields());
+ return (i == time_col_index_) || (i == key_col_index_);
+ }
+
+ // Gets the latest row index, assuming the queue isn't empty
+ row_index_t get_latest_row() const { return latest_ref_row_; }
+
+ bool empty() const {
+ if (latest_ref_row_ > 0)
+ return false; // cannot be empty if ref row is >0 -- can avoid slow queue lock
+ // below
+ return queue_.empty();
+ }
+
+ int countbatches_processed_() const { return batches_processed_; }
+ int count_total_batches() const { return total_batches_; }
+
+ // Gets latest batch (precondition: must not be empty)
+ const std::shared_ptr<arrow::RecordBatch>& get_latest_batch() const {
+ return queue_.unsync_front();
+ }
+ KeyType get_latest_key() const {
+ return queue_.unsync_front()
+ ->column_data(key_col_index_)
+ ->GetValues<KeyType>(1)[latest_ref_row_];
+ }
+ int64_t get_latest_time() const {
+ return queue_.unsync_front()
+ ->column_data(time_col_index_)
+ ->GetValues<int64_t>(1)[latest_ref_row_];
+ }
+
+ bool finished() const { return batches_processed_ == total_batches_; }
+
+ bool advance() {
+ // 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();
+ if (have_active_batch) {
+ // If we have an active batch
+ if (++latest_ref_row_ >= (row_index_t)queue_.unsync_front()->num_rows()) {
+ // hit the end of the batch, need to get the next batch if possible.
+ ++batches_processed_;
+ latest_ref_row_ = 0;
+ have_active_batch &= !queue_.try_pop();
+ if (have_active_batch)
+ assert(queue_.unsync_front()->num_rows() > 0); // empty batches disallowed
+ }
+ }
+ return have_active_batch;
+ }
+
+ // Advance the data to be immediately past the specified TS, updating latest and
+ // latest_ref_row to the latest data prior to that immediate just past Returns true if
+ // updates were made, false if not.
+ bool advance_and_memoize(int64_t ts) {
+ // Advance the right side row index until we reach the latest right row (for each key)
+ // for the given left timestamp.
+
+ // Check if already updated for TS (or if there is no latest)
+ if (empty()) return false; // can't advance if empty
+ auto latest_time = get_latest_time();
+ if (latest_time > ts) return false; // already advanced
+
+ // Not updated. Try to update and possibly advance.
+ bool updated = false;
+ do {
+ latest_time = get_latest_time();
+ // 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 <= ts) {
+ memo_.store(get_latest_batch(), latest_ref_row_, latest_time, get_latest_key());
+ } else {
+ break; // hit a future timestamp -- done updating for now
+ }
+ updated = true;
+ } while (advance());
+ return updated;
+ }
+
+ void push(const std::shared_ptr<arrow::RecordBatch>& rb) {
+ if (rb->num_rows() > 0) {
+ queue_.push(rb);
+ } else {
+ ++batches_processed_; // don't enqueue empty batches, just record as processed
+ }
+ }
+
+ util::optional<const MemoStore::Entry*> get_memo_entry_for_key(KeyType key) {
+ auto r = memo_.get_entry_for_key(key);
+ if (r.has_value()) return r;
+ if (wildcard_key_.has_value()) r = memo_.get_entry_for_key(*wildcard_key_);
+ return r;
+ }
+
+ util::optional<int64_t> get_memo_time_for_key(KeyType key) {
+ auto r = get_memo_entry_for_key(key);
+ return r.has_value() ? util::make_optional((*r)->_time) : util::nullopt;
+ }
+
+ void remove_memo_entries_with_lesser_time(int64_t ts) {
+ memo_.remove_entries_with_lesser_time(ts);
+ }
+
+ const std::shared_ptr<Schema>& get_schema() const { return schema_; }
+
+ void set_total_batches(int n) {
+ assert(n >= 0);
+ assert(total_batches_ == -1); // shouldn't be set more than once
+ total_batches_ = n;
+ }
+
+ private:
+ // Pending record batches. The latest is the front. Batches cannot be empty.
+ ConcurrentQueue<std::shared_ptr<RecordBatch>> queue_;
+
+ // Wildcard key for this input, if applicable.
+ util::optional<KeyType> wildcard_key_;
+
+ // Schema associated with the input
+ std::shared_ptr<Schema> schema_;
+
+ // Total number of batches (only int because InputFinished uses int)
+ int total_batches_ = -1;
+
+ // Number of batches processed so far (only int because InputFinished uses int)
+ int batches_processed_ = 0;
+
+ // Index of the time col
+ col_index_t time_col_index_;
+
+ // Index of the key col
+ col_index_t key_col_index_;
+
+ // Index of the latest row reference within; if >0 then queue_ cannot be empty
+ row_index_t latest_ref_row_ =
+ 0; // must be < queue_.front()->num_rows() if queue_ is non-empty
+
+ // Stores latest known values for the various keys
+ MemoStore memo_;
+
+ // Mapping of source columns to destination columns
+ std::vector<util::optional<col_index_t>> src_to_dst_;
+};
+
+template <size_t MAX_TABLES>
+struct CompositeReferenceRow {
+ struct Entry {
+ arrow::RecordBatch* batch; // can be NULL if there's no value
+ row_index_t row;
+ };
+ Entry refs[MAX_TABLES];
+};
+
+// A table of composite reference rows. Rows maintain pointers to the
+// constituent record batches, but the overall table retains shared_ptr
+// references to ensure memory remains resident while the table is live.
+//
+// The main reason for this is that, especially for wide tables, joins
+// are effectively row-oriented, rather than column-oriented. Separating
+// the join part from the columnar materialization part simplifies the
+// logic around data types and increases efficiency.
+//
+// We don't put the shared_ptr's into the rows for efficiency reasons.
+template <size_t MAX_TABLES>
+class CompositeReferenceTable {
+ public:
+ explicit CompositeReferenceTable(size_t n_tables) : n_tables_(n_tables) {
+ assert(n_tables_ >= 1);
+ assert(n_tables_ <= MAX_TABLES);
+ }
+
+ size_t n_rows() const { return rows_.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<std::unique_ptr<InputState>>& in, int64_t tolerance) {
+ assert(in.size() == n_tables_);
+
+ // Get the LHS key
+ KeyType key = in[0]->get_latest_key();
+
+ // Add row and setup LHS
+ // (the LHS state comes just from the latest row of the LHS table)
+ assert(!in[0]->empty());
+ const std::shared_ptr<arrow::RecordBatch>& lhs_latest_batch =
+ in[0]->get_latest_batch();
+ row_index_t lhs_latest_row = in[0]->get_latest_row();
+ int64_t lhs_latest_time = in[0]->get_latest_time();
+ 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 = rows_.size() + new_batch_size;
+ if (rows_.capacity() < new_capacity) rows_.reserve(new_capacity);
+ }
+ rows_.resize(rows_.size() + 1);
+ auto& row = rows_.back();
+ row.refs[0].batch = lhs_latest_batch.get();
+ row.refs[0].row = lhs_latest_row;
+ add_record_batch_ref(lhs_latest_batch);
+
+ // 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) {
+ util::optional<const MemoStore::Entry*> opt_entry =
+ in[i]->get_memo_entry_for_key(key);
+ if (opt_entry.has_value()) {
+ assert(*opt_entry);
+ if ((*opt_entry)->_time + tolerance >= lhs_latest_time) {
+ // Have a valid entry
+ const MemoStore::Entry* entry = *opt_entry;
+ row.refs[i].batch = entry->_batch.get();
+ row.refs[i].row = entry->_row;
+ add_record_batch_ref(entry->_batch);
+ continue;
+ }
+ }
+ row.refs[i].batch = NULL;
+ row.refs[i].row = 0;
+ }
+ }
+
+ // Materializes the current reference table into a target record batch
+ Result<std::shared_ptr<RecordBatch>> materialize(
+ const std::shared_ptr<arrow::Schema>& output_schema,
+ const std::vector<std::unique_ptr<InputState>>& state) {
+ // cerr << "materialize BEGIN\n";
+ assert(state.size() == n_tables_);
+ assert(state.size() >= 1);
+
+ // Don't build empty batches
+ size_t n_rows = rows_.size();
+ if (!n_rows) return NULLPTR;
+
+ // Build the arrays column-by-column from the rows
+ std::vector<std::shared_ptr<arrow::Array>> arrays(output_schema->num_fields());
+ for (size_t i_table = 0; i_table < n_tables_; ++i_table) {
+ int n_src_cols = state.at(i_table)->get_schema()->num_fields();
+ {
+ for (col_index_t i_src_col = 0; i_src_col < n_src_cols; ++i_src_col) {
+ util::optional<col_index_t> i_dst_col_opt =
+ state[i_table]->map_src_to_dst(i_src_col);
+ if (!i_dst_col_opt) continue;
+ col_index_t i_dst_col = *i_dst_col_opt;
+ const auto& src_field = state[i_table]->get_schema()->field(i_src_col);
+ const auto& dst_field = output_schema->field(i_dst_col);
+ assert(src_field->type()->Equals(dst_field->type()));
+ assert(src_field->name() == dst_field->name());
+ const auto& field_type = src_field->type();
+
+ if (field_type->Equals(arrow::int32())) {
+ ARROW_ASSIGN_OR_RAISE(
+ arrays.at(i_dst_col),
+ (materialize_primitive_column<arrow::Int32Builder, int32_t>(i_table,
+ i_src_col)));
+ } else if (field_type->Equals(arrow::int64())) {
+ ARROW_ASSIGN_OR_RAISE(
+ arrays.at(i_dst_col),
+ (materialize_primitive_column<arrow::Int64Builder, int64_t>(i_table,
+ i_src_col)));
+ } else if (field_type->Equals(arrow::float64())) {
+ ARROW_ASSIGN_OR_RAISE(
+ arrays.at(i_dst_col),
+ (materialize_primitive_column<arrow::DoubleBuilder, double>(i_table,
+ i_src_col)));
+ } else {
+ ARROW_RETURN_NOT_OK(
+ Status::Invalid("Unsupported data type: ", src_field->name()));
+ }
+ }
+ }
+ }
+
+ // Build the result
+ assert(sizeof(size_t) >= sizeof(int64_t)); // Make takes signed int64_t for num_rows
+
+ // TODO: check n_rows for cast
+ std::shared_ptr<arrow::RecordBatch> r =
+ arrow::RecordBatch::Make(output_schema, (int64_t)n_rows, arrays);
+ return r;
+ }
+
+ // Returns true if there are no rows
+ bool empty() const { return rows_.empty(); }
+
+ private:
+ // Contains shared_ptr refs for all RecordBatches referred to by the contents of rows_
+ std::unordered_map<uintptr_t, std::shared_ptr<RecordBatch>> _ptr2ref;
+
+ // Row table references
+ std::vector<CompositeReferenceRow<MAX_TABLES>> rows_;
+
+ // Total number of tables in the composite table
+ size_t n_tables_;
+
+ // Adds a RecordBatch ref to the mapping, if needed
+ void add_record_batch_ref(const std::shared_ptr<RecordBatch>& ref) {
+ if (!_ptr2ref.count((uintptr_t)ref.get())) _ptr2ref[(uintptr_t)ref.get()] = ref;
Review Comment:
You might be able to do:
```
_ptr2ref.try_emplace((uintptr_t)ref.get(), ref);
```
##########
cpp/src/arrow/compute/exec/asof_join_node.cc:
##########
@@ -0,0 +1,800 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+#include <iostream>
+#include <unordered_map>
+
+#include <arrow/api.h>
+#include <arrow/compute/api.h>
+#include <arrow/util/optional.h>
+#include "arrow/compute/exec/asof_join.h"
+#include "arrow/compute/exec/exec_plan.h"
+#include "arrow/compute/exec/options.h"
+#include "arrow/compute/exec/schema_util.h"
+#include "arrow/compute/exec/util.h"
+#include "arrow/result.h"
+#include "arrow/status.h"
+#include "arrow/util/checked_cast.h"
+#include "arrow/util/future.h"
+#include "arrow/util/make_unique.h"
+
+#include <condition_variable>
+#include <mutex>
+#include <thread>
+
+namespace arrow {
+namespace compute {
+
+/**
+ * Simple implementation for an unbound concurrent queue
+ */
+template <class T>
+class ConcurrentQueue {
+ public:
+ T pop() {
+ std::unique_lock<std::mutex> lock(mutex_);
+ cond_.wait(lock, [&] { return !queue_.empty(); });
+ auto item = queue_.front();
+ queue_.pop();
+ return item;
+ }
+
+ void push(const T& item) {
+ std::unique_lock<std::mutex> lock(mutex_);
+ queue_.push(item);
+ cond_.notify_one();
+ }
+
+ util::optional<T> try_pop() {
+ // Try to pop the oldest value from the queue (or return nullopt if none)
+ std::unique_lock<std::mutex> lock(mutex_);
+ if (queue_.empty()) {
+ return util::nullopt;
+ } else {
+ auto item = queue_.front();
+ queue_.pop();
+ return item;
+ }
+ }
+
+ bool empty() const {
+ std::unique_lock<std::mutex> lock(mutex_);
+ return queue_.empty();
+ }
+
+ // Un-synchronized access to front
+ // For this to be "safe":
+ // 1) the caller logically guarantees that queue is not empty
+ // 2) pop/try_pop cannot be called concurrently with this
+ const T& unsync_front() const { return queue_.front(); }
+
+ private:
+ std::queue<T> queue_;
+ mutable std::mutex mutex_;
+ std::condition_variable cond_;
+};
+
+struct MemoStore {
+ // Stores last known values for all the keys
+
+ struct Entry {
+ // Timestamp associated with the entry
+ int64_t _time;
+
+ // Batch associated with the entry (perf is probably OK for this; batches change
+ // rarely)
+ std::shared_ptr<arrow::RecordBatch> _batch;
+
+ // Row associated with the entry
+ row_index_t _row;
+ };
+
+ std::unordered_map<KeyType, Entry> _entries;
+
+ void store(const std::shared_ptr<RecordBatch>& batch, row_index_t row, int64_t time,
+ KeyType key) {
+ auto& e = _entries[key];
+ // that we can do this assignment optionally, is why we
+ // can get array 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;
+ }
+
+ util::optional<const Entry*> get_entry_for_key(KeyType key) const {
+ auto e = _entries.find(key);
+ if (_entries.end() == e) return util::nullopt;
+ return util::optional<const Entry*>(&e->second);
+ }
+
+ void remove_entries_with_lesser_time(int64_t ts) {
+ size_t dbg_size0 = _entries.size();
+ for (auto e = _entries.begin(); e != _entries.end();)
+ if (e->second._time < ts)
+ e = _entries.erase(e);
+ else
+ ++e;
+ size_t dbg_size1 = _entries.size();
+ if (dbg_size1 < dbg_size0) {
+ // cerr << "Removed " << dbg_size0-dbg_size1 << " memo entries.\n";
+ }
+ }
+};
+
+class InputState {
+ // InputState correponds to an input
+ // Input record batches are queued up in InputState until processed and
+ // turned into output record batches.
+
+ public:
+ InputState(const std::shared_ptr<arrow::Schema>& schema,
+ const std::string& time_col_name, const std::string& key_col_name,
+ util::optional<KeyType> wildcard_key)
+ : queue_(),
+ wildcard_key_(wildcard_key),
+ schema_(schema),
+ time_col_index_(
+ schema->GetFieldIndex(time_col_name)), // TODO: handle missing field name
+ key_col_index_(schema->GetFieldIndex(key_col_name)) {}
+
+ col_index_t init_src_to_dst_mapping(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 && is_time_or_key_column(i)))
+ src_to_dst_[i] = dst_offset++;
+ return dst_offset;
+ }
+
+ const util::optional<col_index_t>& map_src_to_dst(col_index_t src) const {
+ return src_to_dst_[src];
+ }
+
+ bool is_time_or_key_column(col_index_t i) const {
+ assert(i < schema_->num_fields());
+ return (i == time_col_index_) || (i == key_col_index_);
+ }
+
+ // Gets the latest row index, assuming the queue isn't empty
+ row_index_t get_latest_row() const { return latest_ref_row_; }
+
+ bool empty() const {
+ if (latest_ref_row_ > 0)
+ return false; // cannot be empty if ref row is >0 -- can avoid slow queue lock
+ // below
+ return queue_.empty();
+ }
+
+ int countbatches_processed_() const { return batches_processed_; }
+ int count_total_batches() const { return total_batches_; }
+
+ // Gets latest batch (precondition: must not be empty)
+ const std::shared_ptr<arrow::RecordBatch>& get_latest_batch() const {
+ return queue_.unsync_front();
+ }
+ KeyType get_latest_key() const {
+ return queue_.unsync_front()
+ ->column_data(key_col_index_)
+ ->GetValues<KeyType>(1)[latest_ref_row_];
+ }
+ int64_t get_latest_time() const {
+ return queue_.unsync_front()
+ ->column_data(time_col_index_)
+ ->GetValues<int64_t>(1)[latest_ref_row_];
+ }
+
+ bool finished() const { return batches_processed_ == total_batches_; }
+
+ bool advance() {
+ // 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();
+ if (have_active_batch) {
+ // If we have an active batch
+ if (++latest_ref_row_ >= (row_index_t)queue_.unsync_front()->num_rows()) {
+ // hit the end of the batch, need to get the next batch if possible.
+ ++batches_processed_;
+ latest_ref_row_ = 0;
+ have_active_batch &= !queue_.try_pop();
+ if (have_active_batch)
+ assert(queue_.unsync_front()->num_rows() > 0); // empty batches disallowed
+ }
+ }
+ return have_active_batch;
+ }
+
+ // Advance the data to be immediately past the specified TS, updating latest and
+ // latest_ref_row to the latest data prior to that immediate just past Returns true if
+ // updates were made, false if not.
+ bool advance_and_memoize(int64_t ts) {
+ // Advance the right side row index until we reach the latest right row (for each key)
+ // for the given left timestamp.
+
+ // Check if already updated for TS (or if there is no latest)
+ if (empty()) return false; // can't advance if empty
+ auto latest_time = get_latest_time();
+ if (latest_time > ts) return false; // already advanced
+
+ // Not updated. Try to update and possibly advance.
+ bool updated = false;
+ do {
+ latest_time = get_latest_time();
+ // 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 <= ts) {
+ memo_.store(get_latest_batch(), latest_ref_row_, latest_time, get_latest_key());
+ } else {
+ break; // hit a future timestamp -- done updating for now
+ }
+ updated = true;
+ } while (advance());
+ return updated;
+ }
+
+ void push(const std::shared_ptr<arrow::RecordBatch>& rb) {
+ if (rb->num_rows() > 0) {
+ queue_.push(rb);
+ } else {
+ ++batches_processed_; // don't enqueue empty batches, just record as processed
+ }
+ }
+
+ util::optional<const MemoStore::Entry*> get_memo_entry_for_key(KeyType key) {
+ auto r = memo_.get_entry_for_key(key);
+ if (r.has_value()) return r;
+ if (wildcard_key_.has_value()) r = memo_.get_entry_for_key(*wildcard_key_);
+ return r;
+ }
+
+ util::optional<int64_t> get_memo_time_for_key(KeyType key) {
+ auto r = get_memo_entry_for_key(key);
+ return r.has_value() ? util::make_optional((*r)->_time) : util::nullopt;
+ }
+
+ void remove_memo_entries_with_lesser_time(int64_t ts) {
+ memo_.remove_entries_with_lesser_time(ts);
+ }
+
+ const std::shared_ptr<Schema>& get_schema() const { return schema_; }
+
+ void set_total_batches(int n) {
+ assert(n >= 0);
+ assert(total_batches_ == -1); // shouldn't be set more than once
+ total_batches_ = n;
+ }
+
+ private:
+ // Pending record batches. The latest is the front. Batches cannot be empty.
+ ConcurrentQueue<std::shared_ptr<RecordBatch>> queue_;
+
+ // Wildcard key for this input, if applicable.
+ util::optional<KeyType> wildcard_key_;
+
+ // Schema associated with the input
+ std::shared_ptr<Schema> schema_;
+
+ // Total number of batches (only int because InputFinished uses int)
+ int total_batches_ = -1;
+
+ // Number of batches processed so far (only int because InputFinished uses int)
+ int batches_processed_ = 0;
+
+ // Index of the time col
+ col_index_t time_col_index_;
+
+ // Index of the key col
+ col_index_t key_col_index_;
+
+ // Index of the latest row reference within; if >0 then queue_ cannot be empty
+ row_index_t latest_ref_row_ =
+ 0; // must be < queue_.front()->num_rows() if queue_ is non-empty
+
+ // Stores latest known values for the various keys
+ MemoStore memo_;
+
+ // Mapping of source columns to destination columns
+ std::vector<util::optional<col_index_t>> src_to_dst_;
+};
+
+template <size_t MAX_TABLES>
+struct CompositeReferenceRow {
+ struct Entry {
+ arrow::RecordBatch* batch; // can be NULL if there's no value
+ row_index_t row;
+ };
+ Entry refs[MAX_TABLES];
+};
+
+// A table of composite reference rows. Rows maintain pointers to the
+// constituent record batches, but the overall table retains shared_ptr
+// references to ensure memory remains resident while the table is live.
+//
+// The main reason for this is that, especially for wide tables, joins
+// are effectively row-oriented, rather than column-oriented. Separating
+// the join part from the columnar materialization part simplifies the
+// logic around data types and increases efficiency.
+//
+// We don't put the shared_ptr's into the rows for efficiency reasons.
+template <size_t MAX_TABLES>
+class CompositeReferenceTable {
+ public:
+ explicit CompositeReferenceTable(size_t n_tables) : n_tables_(n_tables) {
+ assert(n_tables_ >= 1);
+ assert(n_tables_ <= MAX_TABLES);
+ }
+
+ size_t n_rows() const { return rows_.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<std::unique_ptr<InputState>>& in, int64_t tolerance) {
+ assert(in.size() == n_tables_);
+
+ // Get the LHS key
+ KeyType key = in[0]->get_latest_key();
+
+ // Add row and setup LHS
+ // (the LHS state comes just from the latest row of the LHS table)
+ assert(!in[0]->empty());
Review Comment:
```suggestion
DCHECK(!in[0]->empty());
```
##########
cpp/src/arrow/compute/exec/asof_join_node.cc:
##########
@@ -0,0 +1,800 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+#include <iostream>
+#include <unordered_map>
+
+#include <arrow/api.h>
+#include <arrow/compute/api.h>
+#include <arrow/util/optional.h>
+#include "arrow/compute/exec/asof_join.h"
+#include "arrow/compute/exec/exec_plan.h"
+#include "arrow/compute/exec/options.h"
+#include "arrow/compute/exec/schema_util.h"
+#include "arrow/compute/exec/util.h"
+#include "arrow/result.h"
+#include "arrow/status.h"
+#include "arrow/util/checked_cast.h"
+#include "arrow/util/future.h"
+#include "arrow/util/make_unique.h"
+
+#include <condition_variable>
+#include <mutex>
+#include <thread>
+
+namespace arrow {
+namespace compute {
+
+/**
+ * Simple implementation for an unbound concurrent queue
+ */
+template <class T>
+class ConcurrentQueue {
+ public:
+ T pop() {
+ std::unique_lock<std::mutex> lock(mutex_);
+ cond_.wait(lock, [&] { return !queue_.empty(); });
+ auto item = queue_.front();
+ queue_.pop();
+ return item;
+ }
+
+ void push(const T& item) {
+ std::unique_lock<std::mutex> lock(mutex_);
+ queue_.push(item);
+ cond_.notify_one();
+ }
+
+ util::optional<T> try_pop() {
+ // Try to pop the oldest value from the queue (or return nullopt if none)
+ std::unique_lock<std::mutex> lock(mutex_);
+ if (queue_.empty()) {
+ return util::nullopt;
+ } else {
+ auto item = queue_.front();
+ queue_.pop();
+ return item;
+ }
+ }
+
+ bool empty() const {
+ std::unique_lock<std::mutex> lock(mutex_);
+ return queue_.empty();
+ }
+
+ // Un-synchronized access to front
+ // For this to be "safe":
+ // 1) the caller logically guarantees that queue is not empty
+ // 2) pop/try_pop cannot be called concurrently with this
+ const T& unsync_front() const { return queue_.front(); }
+
+ private:
+ std::queue<T> queue_;
+ mutable std::mutex mutex_;
+ std::condition_variable cond_;
+};
+
+struct MemoStore {
+ // Stores last known values for all the keys
+
+ struct Entry {
+ // Timestamp associated with the entry
+ int64_t _time;
+
+ // Batch associated with the entry (perf is probably OK for this; batches change
+ // rarely)
+ std::shared_ptr<arrow::RecordBatch> _batch;
+
+ // Row associated with the entry
+ row_index_t _row;
+ };
+
+ std::unordered_map<KeyType, Entry> _entries;
+
+ void store(const std::shared_ptr<RecordBatch>& batch, row_index_t row, int64_t time,
+ KeyType key) {
+ auto& e = _entries[key];
+ // that we can do this assignment optionally, is why we
+ // can get array 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;
+ }
+
+ util::optional<const Entry*> get_entry_for_key(KeyType key) const {
+ auto e = _entries.find(key);
+ if (_entries.end() == e) return util::nullopt;
+ return util::optional<const Entry*>(&e->second);
+ }
+
+ void remove_entries_with_lesser_time(int64_t ts) {
+ size_t dbg_size0 = _entries.size();
+ for (auto e = _entries.begin(); e != _entries.end();)
+ if (e->second._time < ts)
+ e = _entries.erase(e);
+ else
+ ++e;
+ size_t dbg_size1 = _entries.size();
+ if (dbg_size1 < dbg_size0) {
+ // cerr << "Removed " << dbg_size0-dbg_size1 << " memo entries.\n";
+ }
+ }
+};
+
+class InputState {
+ // InputState correponds to an input
+ // Input record batches are queued up in InputState until processed and
+ // turned into output record batches.
+
+ public:
+ InputState(const std::shared_ptr<arrow::Schema>& schema,
+ const std::string& time_col_name, const std::string& key_col_name,
+ util::optional<KeyType> wildcard_key)
+ : queue_(),
+ wildcard_key_(wildcard_key),
+ schema_(schema),
+ time_col_index_(
+ schema->GetFieldIndex(time_col_name)), // TODO: handle missing field name
+ key_col_index_(schema->GetFieldIndex(key_col_name)) {}
+
+ col_index_t init_src_to_dst_mapping(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 && is_time_or_key_column(i)))
+ src_to_dst_[i] = dst_offset++;
+ return dst_offset;
+ }
+
+ const util::optional<col_index_t>& map_src_to_dst(col_index_t src) const {
+ return src_to_dst_[src];
+ }
+
+ bool is_time_or_key_column(col_index_t i) const {
+ assert(i < schema_->num_fields());
+ return (i == time_col_index_) || (i == key_col_index_);
+ }
+
+ // Gets the latest row index, assuming the queue isn't empty
+ row_index_t get_latest_row() const { return latest_ref_row_; }
+
+ bool empty() const {
+ if (latest_ref_row_ > 0)
+ return false; // cannot be empty if ref row is >0 -- can avoid slow queue lock
+ // below
+ return queue_.empty();
+ }
+
+ int countbatches_processed_() const { return batches_processed_; }
+ int count_total_batches() const { return total_batches_; }
+
+ // Gets latest batch (precondition: must not be empty)
+ const std::shared_ptr<arrow::RecordBatch>& get_latest_batch() const {
+ return queue_.unsync_front();
+ }
+ KeyType get_latest_key() const {
+ return queue_.unsync_front()
+ ->column_data(key_col_index_)
+ ->GetValues<KeyType>(1)[latest_ref_row_];
+ }
+ int64_t get_latest_time() const {
+ return queue_.unsync_front()
+ ->column_data(time_col_index_)
+ ->GetValues<int64_t>(1)[latest_ref_row_];
+ }
+
+ bool finished() const { return batches_processed_ == total_batches_; }
+
+ bool advance() {
+ // 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();
+ if (have_active_batch) {
+ // If we have an active batch
+ if (++latest_ref_row_ >= (row_index_t)queue_.unsync_front()->num_rows()) {
+ // hit the end of the batch, need to get the next batch if possible.
+ ++batches_processed_;
+ latest_ref_row_ = 0;
+ have_active_batch &= !queue_.try_pop();
+ if (have_active_batch)
+ assert(queue_.unsync_front()->num_rows() > 0); // empty batches disallowed
+ }
+ }
+ return have_active_batch;
+ }
+
+ // Advance the data to be immediately past the specified TS, updating latest and
+ // latest_ref_row to the latest data prior to that immediate just past Returns true if
+ // updates were made, false if not.
+ bool advance_and_memoize(int64_t ts) {
+ // Advance the right side row index until we reach the latest right row (for each key)
+ // for the given left timestamp.
+
+ // Check if already updated for TS (or if there is no latest)
+ if (empty()) return false; // can't advance if empty
+ auto latest_time = get_latest_time();
+ if (latest_time > ts) return false; // already advanced
+
+ // Not updated. Try to update and possibly advance.
+ bool updated = false;
+ do {
+ latest_time = get_latest_time();
+ // 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 <= ts) {
+ memo_.store(get_latest_batch(), latest_ref_row_, latest_time, get_latest_key());
+ } else {
+ break; // hit a future timestamp -- done updating for now
+ }
+ updated = true;
+ } while (advance());
+ return updated;
+ }
+
+ void push(const std::shared_ptr<arrow::RecordBatch>& rb) {
+ if (rb->num_rows() > 0) {
+ queue_.push(rb);
+ } else {
+ ++batches_processed_; // don't enqueue empty batches, just record as processed
+ }
+ }
+
+ util::optional<const MemoStore::Entry*> get_memo_entry_for_key(KeyType key) {
+ auto r = memo_.get_entry_for_key(key);
+ if (r.has_value()) return r;
+ if (wildcard_key_.has_value()) r = memo_.get_entry_for_key(*wildcard_key_);
+ return r;
+ }
+
+ util::optional<int64_t> get_memo_time_for_key(KeyType key) {
+ auto r = get_memo_entry_for_key(key);
+ return r.has_value() ? util::make_optional((*r)->_time) : util::nullopt;
+ }
+
+ void remove_memo_entries_with_lesser_time(int64_t ts) {
+ memo_.remove_entries_with_lesser_time(ts);
+ }
+
+ const std::shared_ptr<Schema>& get_schema() const { return schema_; }
+
+ void set_total_batches(int n) {
+ assert(n >= 0);
+ assert(total_batches_ == -1); // shouldn't be set more than once
+ total_batches_ = n;
+ }
+
+ private:
+ // Pending record batches. The latest is the front. Batches cannot be empty.
+ ConcurrentQueue<std::shared_ptr<RecordBatch>> queue_;
+
+ // Wildcard key for this input, if applicable.
+ util::optional<KeyType> wildcard_key_;
+
+ // Schema associated with the input
+ std::shared_ptr<Schema> schema_;
+
+ // Total number of batches (only int because InputFinished uses int)
+ int total_batches_ = -1;
+
+ // Number of batches processed so far (only int because InputFinished uses int)
+ int batches_processed_ = 0;
+
+ // Index of the time col
+ col_index_t time_col_index_;
+
+ // Index of the key col
+ col_index_t key_col_index_;
+
+ // Index of the latest row reference within; if >0 then queue_ cannot be empty
+ row_index_t latest_ref_row_ =
+ 0; // must be < queue_.front()->num_rows() if queue_ is non-empty
+
+ // Stores latest known values for the various keys
+ MemoStore memo_;
+
+ // Mapping of source columns to destination columns
+ std::vector<util::optional<col_index_t>> src_to_dst_;
+};
+
+template <size_t MAX_TABLES>
+struct CompositeReferenceRow {
+ struct Entry {
+ arrow::RecordBatch* batch; // can be NULL if there's no value
+ row_index_t row;
+ };
+ Entry refs[MAX_TABLES];
+};
+
+// A table of composite reference rows. Rows maintain pointers to the
+// constituent record batches, but the overall table retains shared_ptr
+// references to ensure memory remains resident while the table is live.
+//
+// The main reason for this is that, especially for wide tables, joins
+// are effectively row-oriented, rather than column-oriented. Separating
+// the join part from the columnar materialization part simplifies the
+// logic around data types and increases efficiency.
+//
+// We don't put the shared_ptr's into the rows for efficiency reasons.
+template <size_t MAX_TABLES>
+class CompositeReferenceTable {
+ public:
+ explicit CompositeReferenceTable(size_t n_tables) : n_tables_(n_tables) {
+ assert(n_tables_ >= 1);
+ assert(n_tables_ <= MAX_TABLES);
+ }
+
+ size_t n_rows() const { return rows_.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<std::unique_ptr<InputState>>& in, int64_t tolerance) {
+ assert(in.size() == n_tables_);
+
+ // Get the LHS key
+ KeyType key = in[0]->get_latest_key();
+
+ // Add row and setup LHS
+ // (the LHS state comes just from the latest row of the LHS table)
+ assert(!in[0]->empty());
+ const std::shared_ptr<arrow::RecordBatch>& lhs_latest_batch =
+ in[0]->get_latest_batch();
+ row_index_t lhs_latest_row = in[0]->get_latest_row();
+ int64_t lhs_latest_time = in[0]->get_latest_time();
+ 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 = rows_.size() + new_batch_size;
+ if (rows_.capacity() < new_capacity) rows_.reserve(new_capacity);
+ }
+ rows_.resize(rows_.size() + 1);
+ auto& row = rows_.back();
+ row.refs[0].batch = lhs_latest_batch.get();
+ row.refs[0].row = lhs_latest_row;
+ add_record_batch_ref(lhs_latest_batch);
+
+ // 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) {
+ util::optional<const MemoStore::Entry*> opt_entry =
+ in[i]->get_memo_entry_for_key(key);
+ if (opt_entry.has_value()) {
+ assert(*opt_entry);
+ if ((*opt_entry)->_time + tolerance >= lhs_latest_time) {
+ // Have a valid entry
+ const MemoStore::Entry* entry = *opt_entry;
+ row.refs[i].batch = entry->_batch.get();
+ row.refs[i].row = entry->_row;
+ add_record_batch_ref(entry->_batch);
+ continue;
+ }
+ }
+ row.refs[i].batch = NULL;
+ row.refs[i].row = 0;
+ }
+ }
+
+ // Materializes the current reference table into a target record batch
+ Result<std::shared_ptr<RecordBatch>> materialize(
+ const std::shared_ptr<arrow::Schema>& output_schema,
+ const std::vector<std::unique_ptr<InputState>>& state) {
+ // cerr << "materialize BEGIN\n";
+ assert(state.size() == n_tables_);
+ assert(state.size() >= 1);
+
+ // Don't build empty batches
+ size_t n_rows = rows_.size();
+ if (!n_rows) return NULLPTR;
+
+ // Build the arrays column-by-column from the rows
+ std::vector<std::shared_ptr<arrow::Array>> arrays(output_schema->num_fields());
+ for (size_t i_table = 0; i_table < n_tables_; ++i_table) {
+ int n_src_cols = state.at(i_table)->get_schema()->num_fields();
+ {
+ for (col_index_t i_src_col = 0; i_src_col < n_src_cols; ++i_src_col) {
+ util::optional<col_index_t> i_dst_col_opt =
+ state[i_table]->map_src_to_dst(i_src_col);
+ if (!i_dst_col_opt) continue;
+ col_index_t i_dst_col = *i_dst_col_opt;
+ const auto& src_field = state[i_table]->get_schema()->field(i_src_col);
+ const auto& dst_field = output_schema->field(i_dst_col);
+ assert(src_field->type()->Equals(dst_field->type()));
+ assert(src_field->name() == dst_field->name());
+ const auto& field_type = src_field->type();
+
+ if (field_type->Equals(arrow::int32())) {
+ ARROW_ASSIGN_OR_RAISE(
+ arrays.at(i_dst_col),
+ (materialize_primitive_column<arrow::Int32Builder, int32_t>(i_table,
+ i_src_col)));
+ } else if (field_type->Equals(arrow::int64())) {
+ ARROW_ASSIGN_OR_RAISE(
+ arrays.at(i_dst_col),
+ (materialize_primitive_column<arrow::Int64Builder, int64_t>(i_table,
+ i_src_col)));
+ } else if (field_type->Equals(arrow::float64())) {
+ ARROW_ASSIGN_OR_RAISE(
+ arrays.at(i_dst_col),
+ (materialize_primitive_column<arrow::DoubleBuilder, double>(i_table,
+ i_src_col)));
+ } else {
+ ARROW_RETURN_NOT_OK(
+ Status::Invalid("Unsupported data type: ", src_field->name()));
+ }
+ }
+ }
+ }
+
+ // Build the result
+ assert(sizeof(size_t) >= sizeof(int64_t)); // Make takes signed int64_t for num_rows
+
+ // TODO: check n_rows for cast
+ std::shared_ptr<arrow::RecordBatch> r =
+ arrow::RecordBatch::Make(output_schema, (int64_t)n_rows, arrays);
+ return r;
+ }
+
+ // Returns true if there are no rows
+ bool empty() const { return rows_.empty(); }
+
+ private:
+ // Contains shared_ptr refs for all RecordBatches referred to by the contents of rows_
+ std::unordered_map<uintptr_t, std::shared_ptr<RecordBatch>> _ptr2ref;
+
+ // Row table references
+ std::vector<CompositeReferenceRow<MAX_TABLES>> rows_;
+
+ // Total number of tables in the composite table
+ size_t n_tables_;
+
+ // Adds a RecordBatch ref to the mapping, if needed
+ void add_record_batch_ref(const std::shared_ptr<RecordBatch>& ref) {
+ if (!_ptr2ref.count((uintptr_t)ref.get())) _ptr2ref[(uintptr_t)ref.get()] = ref;
+ }
+
+ template <class Builder, class PrimitiveType>
+ Result<std::shared_ptr<Array>> materialize_primitive_column(size_t i_table,
+ col_index_t i_col) {
+ Builder builder;
+ ARROW_RETURN_NOT_OK(builder.Reserve(rows_.size()));
+ for (row_index_t i_row = 0; i_row < rows_.size(); ++i_row) {
+ const auto& ref = rows_[i_row].refs[i_table];
+ if (ref.batch) {
+ builder.UnsafeAppend(
+ ref.batch->column_data(i_col)->template GetValues<PrimitiveType>(1)[ref.row]);
+ } else {
+ builder.UnsafeAppendNull();
+ }
+ }
+ std::shared_ptr<Array> result;
+ ARROW_RETURN_NOT_OK(builder.Finish(&result));
+ return result;
+ }
+};
+
+class AsofJoinNode : public ExecNode {
+ // Constructs labels for inputs
+ static std::vector<std::string> build_input_labels(
+ const std::vector<ExecNode*>& inputs) {
+ std::vector<std::string> r(inputs.size());
+ for (size_t i = 0; i < r.size(); ++i) r[i] = "input_" + std::to_string(i) + "_label";
+ return r;
+ }
+
+ // Advances the RHS as far as possible to be up to date for the current LHS timestamp
+ bool update_rhs() {
+ auto& lhs = *_state.at(0);
+ auto lhs_latest_time = lhs.get_latest_time();
+ bool any_updated = false;
+ for (size_t i = 1; i < _state.size(); ++i)
+ any_updated |= _state[i]->advance_and_memoize(lhs_latest_time);
+ return any_updated;
+ }
+
+ // Returns false if RHS not up to date for LHS
+ bool is_up_to_date_for_lhs_row() const {
+ auto& lhs = *_state[0];
+ if (lhs.empty()) return false; // can't proceed if nothing on the LHS
+ int64_t lhs_ts = lhs.get_latest_time();
+ for (size_t i = 1; i < _state.size(); ++i) {
+ auto& rhs = *_state[i];
+ if (!rhs.finished()) {
+ // If RHS is finished, then we know it's up to date (but if it isn't, it might be
+ // up to date)
+ if (rhs.empty())
+ return false; // RHS isn't finished, but is empty --> not up to date
+ if (lhs_ts >= rhs.get_latest_time())
+ return false; // TS not up to date (and not finished)
+ }
+ }
+ return true;
+ }
+
+ Result<std::shared_ptr<RecordBatch>> process_inner() {
+ assert(!_state.empty());
+ auto& lhs = *_state.at(0);
+
+ // Construct new target table if needed
+ CompositeReferenceTable<MAX_JOIN_TABLES> dst(_state.size());
+
+ // 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;
+
+ // Advance each of the RHS as far as possible to be up to date for the LHS timestamp
+ bool any_advanced = update_rhs();
+
+ // Only update if we have up-to-date information for the LHS row
+ if (is_up_to_date_for_lhs_row()) {
+ dst.emplace(_state, _options.tolerance);
+ if (!lhs.advance()) break; // if we can't advance LHS, we're done for this batch
+ } else {
+ if ((!any_advanced) && (_state.size() > 1)) break; // need to wait for new data
Review Comment:
Why `_state.size() > 1`? What does it mean if `_state.size() == 1`?
##########
cpp/src/arrow/compute/exec/asof_join_node.cc:
##########
@@ -0,0 +1,800 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+#include <iostream>
+#include <unordered_map>
+
+#include <arrow/api.h>
+#include <arrow/compute/api.h>
+#include <arrow/util/optional.h>
+#include "arrow/compute/exec/asof_join.h"
+#include "arrow/compute/exec/exec_plan.h"
+#include "arrow/compute/exec/options.h"
+#include "arrow/compute/exec/schema_util.h"
+#include "arrow/compute/exec/util.h"
+#include "arrow/result.h"
+#include "arrow/status.h"
+#include "arrow/util/checked_cast.h"
+#include "arrow/util/future.h"
+#include "arrow/util/make_unique.h"
+
+#include <condition_variable>
+#include <mutex>
+#include <thread>
+
+namespace arrow {
+namespace compute {
+
+/**
+ * Simple implementation for an unbound concurrent queue
+ */
+template <class T>
+class ConcurrentQueue {
+ public:
+ T pop() {
+ std::unique_lock<std::mutex> lock(mutex_);
+ cond_.wait(lock, [&] { return !queue_.empty(); });
+ auto item = queue_.front();
+ queue_.pop();
+ return item;
+ }
+
+ void push(const T& item) {
+ std::unique_lock<std::mutex> lock(mutex_);
+ queue_.push(item);
+ cond_.notify_one();
+ }
+
+ util::optional<T> try_pop() {
+ // Try to pop the oldest value from the queue (or return nullopt if none)
+ std::unique_lock<std::mutex> lock(mutex_);
+ if (queue_.empty()) {
+ return util::nullopt;
+ } else {
+ auto item = queue_.front();
+ queue_.pop();
+ return item;
+ }
+ }
+
+ bool empty() const {
+ std::unique_lock<std::mutex> lock(mutex_);
+ return queue_.empty();
+ }
+
+ // Un-synchronized access to front
+ // For this to be "safe":
+ // 1) the caller logically guarantees that queue is not empty
+ // 2) pop/try_pop cannot be called concurrently with this
+ const T& unsync_front() const { return queue_.front(); }
+
+ private:
+ std::queue<T> queue_;
+ mutable std::mutex mutex_;
+ std::condition_variable cond_;
+};
+
+struct MemoStore {
+ // Stores last known values for all the keys
+
+ struct Entry {
+ // Timestamp associated with the entry
+ int64_t _time;
+
+ // Batch associated with the entry (perf is probably OK for this; batches change
+ // rarely)
+ std::shared_ptr<arrow::RecordBatch> _batch;
+
+ // Row associated with the entry
+ row_index_t _row;
+ };
+
+ std::unordered_map<KeyType, Entry> _entries;
+
+ void store(const std::shared_ptr<RecordBatch>& batch, row_index_t row, int64_t time,
+ KeyType key) {
+ auto& e = _entries[key];
+ // that we can do this assignment optionally, is why we
+ // can get array 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;
+ }
+
+ util::optional<const Entry*> get_entry_for_key(KeyType key) const {
+ auto e = _entries.find(key);
+ if (_entries.end() == e) return util::nullopt;
+ return util::optional<const Entry*>(&e->second);
+ }
+
+ void remove_entries_with_lesser_time(int64_t ts) {
+ size_t dbg_size0 = _entries.size();
+ for (auto e = _entries.begin(); e != _entries.end();)
+ if (e->second._time < ts)
+ e = _entries.erase(e);
+ else
+ ++e;
+ size_t dbg_size1 = _entries.size();
+ if (dbg_size1 < dbg_size0) {
+ // cerr << "Removed " << dbg_size0-dbg_size1 << " memo entries.\n";
+ }
+ }
+};
+
+class InputState {
+ // InputState correponds to an input
+ // Input record batches are queued up in InputState until processed and
+ // turned into output record batches.
+
+ public:
+ InputState(const std::shared_ptr<arrow::Schema>& schema,
+ const std::string& time_col_name, const std::string& key_col_name,
+ util::optional<KeyType> wildcard_key)
+ : queue_(),
+ wildcard_key_(wildcard_key),
+ schema_(schema),
+ time_col_index_(
+ schema->GetFieldIndex(time_col_name)), // TODO: handle missing field name
+ key_col_index_(schema->GetFieldIndex(key_col_name)) {}
+
+ col_index_t init_src_to_dst_mapping(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 && is_time_or_key_column(i)))
+ src_to_dst_[i] = dst_offset++;
+ return dst_offset;
+ }
+
+ const util::optional<col_index_t>& map_src_to_dst(col_index_t src) const {
+ return src_to_dst_[src];
+ }
+
+ bool is_time_or_key_column(col_index_t i) const {
+ assert(i < schema_->num_fields());
+ return (i == time_col_index_) || (i == key_col_index_);
+ }
+
+ // Gets the latest row index, assuming the queue isn't empty
+ row_index_t get_latest_row() const { return latest_ref_row_; }
+
+ bool empty() const {
+ if (latest_ref_row_ > 0)
+ return false; // cannot be empty if ref row is >0 -- can avoid slow queue lock
+ // below
+ return queue_.empty();
+ }
+
+ int countbatches_processed_() const { return batches_processed_; }
+ int count_total_batches() const { return total_batches_; }
+
+ // Gets latest batch (precondition: must not be empty)
+ const std::shared_ptr<arrow::RecordBatch>& get_latest_batch() const {
+ return queue_.unsync_front();
+ }
+ KeyType get_latest_key() const {
+ return queue_.unsync_front()
+ ->column_data(key_col_index_)
+ ->GetValues<KeyType>(1)[latest_ref_row_];
+ }
+ int64_t get_latest_time() const {
+ return queue_.unsync_front()
+ ->column_data(time_col_index_)
+ ->GetValues<int64_t>(1)[latest_ref_row_];
+ }
+
+ bool finished() const { return batches_processed_ == total_batches_; }
+
+ bool advance() {
+ // 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();
+ if (have_active_batch) {
+ // If we have an active batch
+ if (++latest_ref_row_ >= (row_index_t)queue_.unsync_front()->num_rows()) {
+ // hit the end of the batch, need to get the next batch if possible.
+ ++batches_processed_;
+ latest_ref_row_ = 0;
+ have_active_batch &= !queue_.try_pop();
+ if (have_active_batch)
+ assert(queue_.unsync_front()->num_rows() > 0); // empty batches disallowed
+ }
+ }
+ return have_active_batch;
+ }
+
+ // Advance the data to be immediately past the specified TS, updating latest and
+ // latest_ref_row to the latest data prior to that immediate just past Returns true if
+ // updates were made, false if not.
+ bool advance_and_memoize(int64_t ts) {
+ // Advance the right side row index until we reach the latest right row (for each key)
+ // for the given left timestamp.
+
+ // Check if already updated for TS (or if there is no latest)
+ if (empty()) return false; // can't advance if empty
+ auto latest_time = get_latest_time();
+ if (latest_time > ts) return false; // already advanced
+
+ // Not updated. Try to update and possibly advance.
+ bool updated = false;
+ do {
+ latest_time = get_latest_time();
+ // 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 <= ts) {
+ memo_.store(get_latest_batch(), latest_ref_row_, latest_time, get_latest_key());
+ } else {
+ break; // hit a future timestamp -- done updating for now
+ }
+ updated = true;
+ } while (advance());
+ return updated;
+ }
+
+ void push(const std::shared_ptr<arrow::RecordBatch>& rb) {
+ if (rb->num_rows() > 0) {
+ queue_.push(rb);
+ } else {
+ ++batches_processed_; // don't enqueue empty batches, just record as processed
+ }
+ }
+
+ util::optional<const MemoStore::Entry*> get_memo_entry_for_key(KeyType key) {
+ auto r = memo_.get_entry_for_key(key);
+ if (r.has_value()) return r;
+ if (wildcard_key_.has_value()) r = memo_.get_entry_for_key(*wildcard_key_);
+ return r;
+ }
+
+ util::optional<int64_t> get_memo_time_for_key(KeyType key) {
+ auto r = get_memo_entry_for_key(key);
+ return r.has_value() ? util::make_optional((*r)->_time) : util::nullopt;
+ }
+
+ void remove_memo_entries_with_lesser_time(int64_t ts) {
+ memo_.remove_entries_with_lesser_time(ts);
+ }
+
+ const std::shared_ptr<Schema>& get_schema() const { return schema_; }
+
+ void set_total_batches(int n) {
+ assert(n >= 0);
+ assert(total_batches_ == -1); // shouldn't be set more than once
+ total_batches_ = n;
+ }
+
+ private:
+ // Pending record batches. The latest is the front. Batches cannot be empty.
+ ConcurrentQueue<std::shared_ptr<RecordBatch>> queue_;
+
+ // Wildcard key for this input, if applicable.
+ util::optional<KeyType> wildcard_key_;
+
+ // Schema associated with the input
+ std::shared_ptr<Schema> schema_;
+
+ // Total number of batches (only int because InputFinished uses int)
+ int total_batches_ = -1;
+
+ // Number of batches processed so far (only int because InputFinished uses int)
+ int batches_processed_ = 0;
+
+ // Index of the time col
+ col_index_t time_col_index_;
+
+ // Index of the key col
+ col_index_t key_col_index_;
+
+ // Index of the latest row reference within; if >0 then queue_ cannot be empty
+ row_index_t latest_ref_row_ =
+ 0; // must be < queue_.front()->num_rows() if queue_ is non-empty
+
+ // Stores latest known values for the various keys
+ MemoStore memo_;
+
+ // Mapping of source columns to destination columns
+ std::vector<util::optional<col_index_t>> src_to_dst_;
+};
+
+template <size_t MAX_TABLES>
+struct CompositeReferenceRow {
+ struct Entry {
+ arrow::RecordBatch* batch; // can be NULL if there's no value
+ row_index_t row;
+ };
+ Entry refs[MAX_TABLES];
+};
+
+// A table of composite reference rows. Rows maintain pointers to the
+// constituent record batches, but the overall table retains shared_ptr
+// references to ensure memory remains resident while the table is live.
+//
+// The main reason for this is that, especially for wide tables, joins
+// are effectively row-oriented, rather than column-oriented. Separating
+// the join part from the columnar materialization part simplifies the
+// logic around data types and increases efficiency.
+//
+// We don't put the shared_ptr's into the rows for efficiency reasons.
+template <size_t MAX_TABLES>
+class CompositeReferenceTable {
+ public:
+ explicit CompositeReferenceTable(size_t n_tables) : n_tables_(n_tables) {
+ assert(n_tables_ >= 1);
+ assert(n_tables_ <= MAX_TABLES);
+ }
+
+ size_t n_rows() const { return rows_.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<std::unique_ptr<InputState>>& in, int64_t tolerance) {
+ assert(in.size() == n_tables_);
+
+ // Get the LHS key
+ KeyType key = in[0]->get_latest_key();
+
+ // Add row and setup LHS
+ // (the LHS state comes just from the latest row of the LHS table)
+ assert(!in[0]->empty());
+ const std::shared_ptr<arrow::RecordBatch>& lhs_latest_batch =
+ in[0]->get_latest_batch();
+ row_index_t lhs_latest_row = in[0]->get_latest_row();
+ int64_t lhs_latest_time = in[0]->get_latest_time();
+ 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 = rows_.size() + new_batch_size;
+ if (rows_.capacity() < new_capacity) rows_.reserve(new_capacity);
+ }
+ rows_.resize(rows_.size() + 1);
+ auto& row = rows_.back();
+ row.refs[0].batch = lhs_latest_batch.get();
+ row.refs[0].row = lhs_latest_row;
+ add_record_batch_ref(lhs_latest_batch);
+
+ // 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) {
+ util::optional<const MemoStore::Entry*> opt_entry =
+ in[i]->get_memo_entry_for_key(key);
+ if (opt_entry.has_value()) {
+ assert(*opt_entry);
+ if ((*opt_entry)->_time + tolerance >= lhs_latest_time) {
+ // Have a valid entry
+ const MemoStore::Entry* entry = *opt_entry;
+ row.refs[i].batch = entry->_batch.get();
+ row.refs[i].row = entry->_row;
+ add_record_batch_ref(entry->_batch);
+ continue;
+ }
+ }
+ row.refs[i].batch = NULL;
+ row.refs[i].row = 0;
+ }
+ }
+
+ // Materializes the current reference table into a target record batch
+ Result<std::shared_ptr<RecordBatch>> materialize(
+ const std::shared_ptr<arrow::Schema>& output_schema,
+ const std::vector<std::unique_ptr<InputState>>& state) {
+ // cerr << "materialize BEGIN\n";
+ assert(state.size() == n_tables_);
+ assert(state.size() >= 1);
+
+ // Don't build empty batches
+ size_t n_rows = rows_.size();
+ if (!n_rows) return NULLPTR;
+
+ // Build the arrays column-by-column from the rows
+ std::vector<std::shared_ptr<arrow::Array>> arrays(output_schema->num_fields());
+ for (size_t i_table = 0; i_table < n_tables_; ++i_table) {
+ int n_src_cols = state.at(i_table)->get_schema()->num_fields();
+ {
+ for (col_index_t i_src_col = 0; i_src_col < n_src_cols; ++i_src_col) {
+ util::optional<col_index_t> i_dst_col_opt =
+ state[i_table]->map_src_to_dst(i_src_col);
+ if (!i_dst_col_opt) continue;
+ col_index_t i_dst_col = *i_dst_col_opt;
+ const auto& src_field = state[i_table]->get_schema()->field(i_src_col);
+ const auto& dst_field = output_schema->field(i_dst_col);
+ assert(src_field->type()->Equals(dst_field->type()));
+ assert(src_field->name() == dst_field->name());
+ const auto& field_type = src_field->type();
+
+ if (field_type->Equals(arrow::int32())) {
+ ARROW_ASSIGN_OR_RAISE(
+ arrays.at(i_dst_col),
+ (materialize_primitive_column<arrow::Int32Builder, int32_t>(i_table,
+ i_src_col)));
+ } else if (field_type->Equals(arrow::int64())) {
+ ARROW_ASSIGN_OR_RAISE(
+ arrays.at(i_dst_col),
+ (materialize_primitive_column<arrow::Int64Builder, int64_t>(i_table,
+ i_src_col)));
+ } else if (field_type->Equals(arrow::float64())) {
+ ARROW_ASSIGN_OR_RAISE(
+ arrays.at(i_dst_col),
+ (materialize_primitive_column<arrow::DoubleBuilder, double>(i_table,
+ i_src_col)));
+ } else {
+ ARROW_RETURN_NOT_OK(
+ Status::Invalid("Unsupported data type: ", src_field->name()));
+ }
+ }
+ }
+ }
+
+ // Build the result
+ assert(sizeof(size_t) >= sizeof(int64_t)); // Make takes signed int64_t for num_rows
+
+ // TODO: check n_rows for cast
+ std::shared_ptr<arrow::RecordBatch> r =
+ arrow::RecordBatch::Make(output_schema, (int64_t)n_rows, arrays);
+ return r;
+ }
+
+ // Returns true if there are no rows
+ bool empty() const { return rows_.empty(); }
+
+ private:
+ // Contains shared_ptr refs for all RecordBatches referred to by the contents of rows_
+ std::unordered_map<uintptr_t, std::shared_ptr<RecordBatch>> _ptr2ref;
+
+ // Row table references
+ std::vector<CompositeReferenceRow<MAX_TABLES>> rows_;
+
+ // Total number of tables in the composite table
+ size_t n_tables_;
+
+ // Adds a RecordBatch ref to the mapping, if needed
+ void add_record_batch_ref(const std::shared_ptr<RecordBatch>& ref) {
+ if (!_ptr2ref.count((uintptr_t)ref.get())) _ptr2ref[(uintptr_t)ref.get()] = ref;
+ }
+
+ template <class Builder, class PrimitiveType>
+ Result<std::shared_ptr<Array>> materialize_primitive_column(size_t i_table,
+ col_index_t i_col) {
+ Builder builder;
+ ARROW_RETURN_NOT_OK(builder.Reserve(rows_.size()));
+ for (row_index_t i_row = 0; i_row < rows_.size(); ++i_row) {
+ const auto& ref = rows_[i_row].refs[i_table];
+ if (ref.batch) {
+ builder.UnsafeAppend(
+ ref.batch->column_data(i_col)->template GetValues<PrimitiveType>(1)[ref.row]);
+ } else {
+ builder.UnsafeAppendNull();
+ }
+ }
+ std::shared_ptr<Array> result;
+ ARROW_RETURN_NOT_OK(builder.Finish(&result));
+ return result;
+ }
+};
+
+class AsofJoinNode : public ExecNode {
+ // Constructs labels for inputs
+ static std::vector<std::string> build_input_labels(
+ const std::vector<ExecNode*>& inputs) {
+ std::vector<std::string> r(inputs.size());
+ for (size_t i = 0; i < r.size(); ++i) r[i] = "input_" + std::to_string(i) + "_label";
+ return r;
+ }
+
+ // Advances the RHS as far as possible to be up to date for the current LHS timestamp
+ bool update_rhs() {
+ auto& lhs = *_state.at(0);
+ auto lhs_latest_time = lhs.get_latest_time();
+ bool any_updated = false;
+ for (size_t i = 1; i < _state.size(); ++i)
+ any_updated |= _state[i]->advance_and_memoize(lhs_latest_time);
+ return any_updated;
+ }
+
+ // Returns false if RHS not up to date for LHS
+ bool is_up_to_date_for_lhs_row() const {
+ auto& lhs = *_state[0];
+ if (lhs.empty()) return false; // can't proceed if nothing on the LHS
+ int64_t lhs_ts = lhs.get_latest_time();
+ for (size_t i = 1; i < _state.size(); ++i) {
+ auto& rhs = *_state[i];
+ if (!rhs.finished()) {
+ // If RHS is finished, then we know it's up to date (but if it isn't, it might be
+ // up to date)
+ if (rhs.empty())
+ return false; // RHS isn't finished, but is empty --> not up to date
+ if (lhs_ts >= rhs.get_latest_time())
+ return false; // TS not up to date (and not finished)
+ }
+ }
+ return true;
+ }
+
+ Result<std::shared_ptr<RecordBatch>> process_inner() {
+ assert(!_state.empty());
+ auto& lhs = *_state.at(0);
+
+ // Construct new target table if needed
+ CompositeReferenceTable<MAX_JOIN_TABLES> dst(_state.size());
+
+ // 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;
+
+ // Advance each of the RHS as far as possible to be up to date for the LHS timestamp
+ bool any_advanced = update_rhs();
+
+ // Only update if we have up-to-date information for the LHS row
+ if (is_up_to_date_for_lhs_row()) {
+ dst.emplace(_state, _options.tolerance);
+ if (!lhs.advance()) break; // if we can't advance LHS, we're done for this batch
+ } else {
+ if ((!any_advanced) && (_state.size() > 1)) break; // need to wait for new data
+ }
+ }
+
+ // Prune memo entries that have expired (to bound memory consumption)
+ if (!lhs.empty()) {
+ for (size_t i = 1; i < _state.size(); ++i) {
+ _state[i]->remove_memo_entries_with_lesser_time(lhs.get_latest_time() -
+ _options.tolerance);
+ }
+ }
+
+ // Emit the batch
+ if (dst.empty()) {
+ return NULLPTR;
+ } else {
+ return dst.materialize(output_schema(), _state);
+ }
+ }
+
+ void process() {
+ std::cerr << "process() begin\n";
+
+ std::lock_guard<std::mutex> guard(_gate);
+ if (finished_.is_finished()) {
+ std::cerr << "InputReceived EARLYEND\n";
+ return;
+ }
+
+ // Process batches while we have data
+ for (;;) {
+ Result<std::shared_ptr<RecordBatch>> result = process_inner();
+
+ if (result.ok()) {
+ auto out_rb = *result;
+ if (!out_rb) break;
+ ++_progress_batches_produced;
+ ExecBatch out_b(*out_rb);
+ outputs_[0]->InputReceived(this, std::move(out_b));
+ } else {
+ StopProducing();
+ ErrorIfNotOk(result.status());
+ return;
+ }
+ }
+
+ std::cerr << "process() end\n";
+
+ // 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 (_state.at(0)->finished()) {
+ total_batches_produced_ = util::make_optional<int>(_progress_batches_produced);
+ StopProducing();
+ assert(total_batches_produced_.has_value());
+ outputs_[0]->InputFinished(this, *total_batches_produced_);
+ }
+ }
+
+ void process_thread() {
+ std::cerr << "AsofJoinNode::process_thread started.\n";
+ for (;;) {
+ if (!_process.pop()) {
+ std::cerr << "AsofJoinNode::process_thread done.\n";
+ return;
+ }
+ process();
+ }
+ }
+
+ static void process_thread_wrapper(AsofJoinNode* node) { node->process_thread(); }
+
+ public:
+ AsofJoinNode(ExecPlan* plan, NodeVector inputs, std::vector<std::string> input_labels,
+ const AsofJoinNodeOptions& join_options,
+ std::shared_ptr<Schema> output_schema,
+ std::unique_ptr<AsofJoinSchema> schema_mgr);
+
+ virtual ~AsofJoinNode() {
+ _process.push(false); // poison pill
+ _process_thread.join();
+ }
+
+ static arrow::Result<ExecNode*> Make(ExecPlan* plan, std::vector<ExecNode*> inputs,
+ const ExecNodeOptions& options) {
+ std::unique_ptr<AsofJoinSchema> schema_mgr =
+ ::arrow::internal::make_unique<AsofJoinSchema>();
+
+ const auto& join_options = checked_cast<const AsofJoinNodeOptions&>(options);
+ std::shared_ptr<Schema> output_schema =
+ schema_mgr->MakeOutputSchema(inputs, join_options);
+
+ std::vector<std::string> input_labels(inputs.size());
+ input_labels[0] = "left";
+ for (size_t i = 1; i < inputs.size(); ++i) {
+ input_labels[i] = "right_" + std::to_string(i);
+ }
+
+ return plan->EmplaceNode<AsofJoinNode>(plan, inputs, std::move(input_labels),
+ join_options, std::move(output_schema),
+ std::move(schema_mgr));
+ }
+
+ const char* kind_name() const override { return "AsofJoinNode"; }
+
+ void InputReceived(ExecNode* input, ExecBatch batch) override {
+ // Get the input
+ ARROW_DCHECK(std::find(inputs_.begin(), inputs_.end(), input) != inputs_.end());
+ size_t k = std::find(inputs_.begin(), inputs_.end(), input) - inputs_.begin();
+ std::cerr << "InputReceived BEGIN (k=" << k << ")\n";
+
+ // Put into the queue
+ auto rb = *batch.ToRecordBatch(input->output_schema());
+
+ _state.at(k)->push(rb);
+ _process.push(true);
+
+ std::cerr << "InputReceived END\n";
+ }
+ void ErrorReceived(ExecNode* input, Status error) override {
+ outputs_[0]->ErrorReceived(this, std::move(error));
+ StopProducing();
+ }
+ void InputFinished(ExecNode* input, int total_batches) override {
+ std::cerr << "InputFinished BEGIN\n";
+ // bool is_finished=false;
+ {
+ std::lock_guard<std::mutex> guard(_gate);
+ std::cerr << "InputFinished find\n";
+ ARROW_DCHECK(std::find(inputs_.begin(), inputs_.end(), input) != inputs_.end());
+ size_t k = std::find(inputs_.begin(), inputs_.end(), input) - inputs_.begin();
+ // cerr << "set_total_batches for input " << k << ": " << total_batches << "\n";
+ _state.at(k)->set_total_batches(total_batches);
+ }
+ // 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.
+ _process.push(true);
+
+ std::cerr << "InputFinished END\n";
+ }
+ Status StartProducing() override {
+ std::cout << "StartProducing"
+ << "\n";
+ finished_ = arrow::Future<>::Make();
+ return Status::OK();
+ }
+ void PauseProducing(ExecNode* output, int32_t counter) override {
+ std::cout << "PauseProducing"
+ << "\n";
+ }
+ void ResumeProducing(ExecNode* output, int32_t counter) override {
+ std::cout << "ResumeProducing"
+ << "\n";
+ }
+ void StopProducing(ExecNode* output) override {
+ DCHECK_EQ(output, outputs_[0]);
+ StopProducing();
+ std::cout << "StopProducing"
+ << "\n";
+ }
+ void StopProducing() override {
Review Comment:
Ideally we would stop the process thread here but it will eventually get stopped when the plan is destroyed so maybe not critical.
##########
cpp/src/arrow/compute/exec/asof_join_node.cc:
##########
@@ -0,0 +1,800 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+#include <iostream>
+#include <unordered_map>
+
+#include <arrow/api.h>
+#include <arrow/compute/api.h>
+#include <arrow/util/optional.h>
+#include "arrow/compute/exec/asof_join.h"
+#include "arrow/compute/exec/exec_plan.h"
+#include "arrow/compute/exec/options.h"
+#include "arrow/compute/exec/schema_util.h"
+#include "arrow/compute/exec/util.h"
+#include "arrow/result.h"
+#include "arrow/status.h"
+#include "arrow/util/checked_cast.h"
+#include "arrow/util/future.h"
+#include "arrow/util/make_unique.h"
+
+#include <condition_variable>
+#include <mutex>
+#include <thread>
+
+namespace arrow {
+namespace compute {
+
+/**
+ * Simple implementation for an unbound concurrent queue
+ */
+template <class T>
+class ConcurrentQueue {
+ public:
+ T pop() {
+ std::unique_lock<std::mutex> lock(mutex_);
+ cond_.wait(lock, [&] { return !queue_.empty(); });
+ auto item = queue_.front();
+ queue_.pop();
+ return item;
+ }
+
+ void push(const T& item) {
+ std::unique_lock<std::mutex> lock(mutex_);
+ queue_.push(item);
+ cond_.notify_one();
+ }
+
+ util::optional<T> try_pop() {
+ // Try to pop the oldest value from the queue (or return nullopt if none)
+ std::unique_lock<std::mutex> lock(mutex_);
+ if (queue_.empty()) {
+ return util::nullopt;
+ } else {
+ auto item = queue_.front();
+ queue_.pop();
+ return item;
+ }
+ }
+
+ bool empty() const {
+ std::unique_lock<std::mutex> lock(mutex_);
+ return queue_.empty();
+ }
+
+ // Un-synchronized access to front
+ // For this to be "safe":
+ // 1) the caller logically guarantees that queue is not empty
+ // 2) pop/try_pop cannot be called concurrently with this
+ const T& unsync_front() const { return queue_.front(); }
+
+ private:
+ std::queue<T> queue_;
+ mutable std::mutex mutex_;
+ std::condition_variable cond_;
+};
+
+struct MemoStore {
+ // Stores last known values for all the keys
+
+ struct Entry {
+ // Timestamp associated with the entry
+ int64_t _time;
+
+ // Batch associated with the entry (perf is probably OK for this; batches change
+ // rarely)
+ std::shared_ptr<arrow::RecordBatch> _batch;
+
+ // Row associated with the entry
+ row_index_t _row;
+ };
+
+ std::unordered_map<KeyType, Entry> _entries;
+
+ void store(const std::shared_ptr<RecordBatch>& batch, row_index_t row, int64_t time,
+ KeyType key) {
+ auto& e = _entries[key];
+ // that we can do this assignment optionally, is why we
+ // can get array 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;
+ }
+
+ util::optional<const Entry*> get_entry_for_key(KeyType key) const {
+ auto e = _entries.find(key);
+ if (_entries.end() == e) return util::nullopt;
+ return util::optional<const Entry*>(&e->second);
+ }
+
+ void remove_entries_with_lesser_time(int64_t ts) {
+ size_t dbg_size0 = _entries.size();
+ for (auto e = _entries.begin(); e != _entries.end();)
+ if (e->second._time < ts)
+ e = _entries.erase(e);
+ else
+ ++e;
+ size_t dbg_size1 = _entries.size();
+ if (dbg_size1 < dbg_size0) {
+ // cerr << "Removed " << dbg_size0-dbg_size1 << " memo entries.\n";
+ }
+ }
+};
+
+class InputState {
+ // InputState correponds to an input
+ // Input record batches are queued up in InputState until processed and
+ // turned into output record batches.
+
+ public:
+ InputState(const std::shared_ptr<arrow::Schema>& schema,
+ const std::string& time_col_name, const std::string& key_col_name,
+ util::optional<KeyType> wildcard_key)
+ : queue_(),
+ wildcard_key_(wildcard_key),
+ schema_(schema),
+ time_col_index_(
+ schema->GetFieldIndex(time_col_name)), // TODO: handle missing field name
+ key_col_index_(schema->GetFieldIndex(key_col_name)) {}
+
+ col_index_t init_src_to_dst_mapping(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 && is_time_or_key_column(i)))
+ src_to_dst_[i] = dst_offset++;
+ return dst_offset;
+ }
+
+ const util::optional<col_index_t>& map_src_to_dst(col_index_t src) const {
+ return src_to_dst_[src];
+ }
+
+ bool is_time_or_key_column(col_index_t i) const {
+ assert(i < schema_->num_fields());
+ return (i == time_col_index_) || (i == key_col_index_);
+ }
+
+ // Gets the latest row index, assuming the queue isn't empty
+ row_index_t get_latest_row() const { return latest_ref_row_; }
+
+ bool empty() const {
+ if (latest_ref_row_ > 0)
+ return false; // cannot be empty if ref row is >0 -- can avoid slow queue lock
+ // below
+ return queue_.empty();
+ }
+
+ int countbatches_processed_() const { return batches_processed_; }
+ int count_total_batches() const { return total_batches_; }
+
+ // Gets latest batch (precondition: must not be empty)
+ const std::shared_ptr<arrow::RecordBatch>& get_latest_batch() const {
+ return queue_.unsync_front();
+ }
+ KeyType get_latest_key() const {
+ return queue_.unsync_front()
+ ->column_data(key_col_index_)
+ ->GetValues<KeyType>(1)[latest_ref_row_];
+ }
+ int64_t get_latest_time() const {
+ return queue_.unsync_front()
+ ->column_data(time_col_index_)
+ ->GetValues<int64_t>(1)[latest_ref_row_];
+ }
+
+ bool finished() const { return batches_processed_ == total_batches_; }
+
+ bool advance() {
+ // 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();
+ if (have_active_batch) {
+ // If we have an active batch
+ if (++latest_ref_row_ >= (row_index_t)queue_.unsync_front()->num_rows()) {
+ // hit the end of the batch, need to get the next batch if possible.
+ ++batches_processed_;
+ latest_ref_row_ = 0;
+ have_active_batch &= !queue_.try_pop();
+ if (have_active_batch)
+ assert(queue_.unsync_front()->num_rows() > 0); // empty batches disallowed
+ }
+ }
+ return have_active_batch;
+ }
+
+ // Advance the data to be immediately past the specified TS, updating latest and
+ // latest_ref_row to the latest data prior to that immediate just past Returns true if
+ // updates were made, false if not.
+ bool advance_and_memoize(int64_t ts) {
+ // Advance the right side row index until we reach the latest right row (for each key)
+ // for the given left timestamp.
+
+ // Check if already updated for TS (or if there is no latest)
+ if (empty()) return false; // can't advance if empty
+ auto latest_time = get_latest_time();
+ if (latest_time > ts) return false; // already advanced
+
+ // Not updated. Try to update and possibly advance.
+ bool updated = false;
+ do {
+ latest_time = get_latest_time();
+ // 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 <= ts) {
+ memo_.store(get_latest_batch(), latest_ref_row_, latest_time, get_latest_key());
+ } else {
+ break; // hit a future timestamp -- done updating for now
+ }
+ updated = true;
+ } while (advance());
+ return updated;
+ }
+
+ void push(const std::shared_ptr<arrow::RecordBatch>& rb) {
+ if (rb->num_rows() > 0) {
+ queue_.push(rb);
+ } else {
+ ++batches_processed_; // don't enqueue empty batches, just record as processed
+ }
+ }
+
+ util::optional<const MemoStore::Entry*> get_memo_entry_for_key(KeyType key) {
+ auto r = memo_.get_entry_for_key(key);
+ if (r.has_value()) return r;
+ if (wildcard_key_.has_value()) r = memo_.get_entry_for_key(*wildcard_key_);
+ return r;
+ }
+
+ util::optional<int64_t> get_memo_time_for_key(KeyType key) {
+ auto r = get_memo_entry_for_key(key);
+ return r.has_value() ? util::make_optional((*r)->_time) : util::nullopt;
+ }
+
+ void remove_memo_entries_with_lesser_time(int64_t ts) {
+ memo_.remove_entries_with_lesser_time(ts);
+ }
+
+ const std::shared_ptr<Schema>& get_schema() const { return schema_; }
+
+ void set_total_batches(int n) {
+ assert(n >= 0);
+ assert(total_batches_ == -1); // shouldn't be set more than once
+ total_batches_ = n;
+ }
+
+ private:
+ // Pending record batches. The latest is the front. Batches cannot be empty.
+ ConcurrentQueue<std::shared_ptr<RecordBatch>> queue_;
+
+ // Wildcard key for this input, if applicable.
+ util::optional<KeyType> wildcard_key_;
+
+ // Schema associated with the input
+ std::shared_ptr<Schema> schema_;
+
+ // Total number of batches (only int because InputFinished uses int)
+ int total_batches_ = -1;
+
+ // Number of batches processed so far (only int because InputFinished uses int)
+ int batches_processed_ = 0;
+
+ // Index of the time col
+ col_index_t time_col_index_;
+
+ // Index of the key col
+ col_index_t key_col_index_;
+
+ // Index of the latest row reference within; if >0 then queue_ cannot be empty
+ row_index_t latest_ref_row_ =
+ 0; // must be < queue_.front()->num_rows() if queue_ is non-empty
+
+ // Stores latest known values for the various keys
+ MemoStore memo_;
+
+ // Mapping of source columns to destination columns
+ std::vector<util::optional<col_index_t>> src_to_dst_;
+};
+
+template <size_t MAX_TABLES>
+struct CompositeReferenceRow {
+ struct Entry {
+ arrow::RecordBatch* batch; // can be NULL if there's no value
+ row_index_t row;
+ };
+ Entry refs[MAX_TABLES];
+};
+
+// A table of composite reference rows. Rows maintain pointers to the
+// constituent record batches, but the overall table retains shared_ptr
+// references to ensure memory remains resident while the table is live.
+//
+// The main reason for this is that, especially for wide tables, joins
+// are effectively row-oriented, rather than column-oriented. Separating
+// the join part from the columnar materialization part simplifies the
+// logic around data types and increases efficiency.
+//
+// We don't put the shared_ptr's into the rows for efficiency reasons.
+template <size_t MAX_TABLES>
+class CompositeReferenceTable {
+ public:
+ explicit CompositeReferenceTable(size_t n_tables) : n_tables_(n_tables) {
+ assert(n_tables_ >= 1);
+ assert(n_tables_ <= MAX_TABLES);
+ }
+
+ size_t n_rows() const { return rows_.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<std::unique_ptr<InputState>>& in, int64_t tolerance) {
+ assert(in.size() == n_tables_);
+
+ // Get the LHS key
+ KeyType key = in[0]->get_latest_key();
+
+ // Add row and setup LHS
+ // (the LHS state comes just from the latest row of the LHS table)
+ assert(!in[0]->empty());
+ const std::shared_ptr<arrow::RecordBatch>& lhs_latest_batch =
+ in[0]->get_latest_batch();
+ row_index_t lhs_latest_row = in[0]->get_latest_row();
+ int64_t lhs_latest_time = in[0]->get_latest_time();
+ 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 = rows_.size() + new_batch_size;
+ if (rows_.capacity() < new_capacity) rows_.reserve(new_capacity);
+ }
+ rows_.resize(rows_.size() + 1);
+ auto& row = rows_.back();
+ row.refs[0].batch = lhs_latest_batch.get();
+ row.refs[0].row = lhs_latest_row;
+ add_record_batch_ref(lhs_latest_batch);
+
+ // 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) {
+ util::optional<const MemoStore::Entry*> opt_entry =
+ in[i]->get_memo_entry_for_key(key);
+ if (opt_entry.has_value()) {
+ assert(*opt_entry);
+ if ((*opt_entry)->_time + tolerance >= lhs_latest_time) {
+ // Have a valid entry
+ const MemoStore::Entry* entry = *opt_entry;
+ row.refs[i].batch = entry->_batch.get();
+ row.refs[i].row = entry->_row;
+ add_record_batch_ref(entry->_batch);
+ continue;
+ }
+ }
+ row.refs[i].batch = NULL;
+ row.refs[i].row = 0;
+ }
+ }
+
+ // Materializes the current reference table into a target record batch
+ Result<std::shared_ptr<RecordBatch>> materialize(
+ const std::shared_ptr<arrow::Schema>& output_schema,
+ const std::vector<std::unique_ptr<InputState>>& state) {
+ // cerr << "materialize BEGIN\n";
+ assert(state.size() == n_tables_);
+ assert(state.size() >= 1);
+
+ // Don't build empty batches
+ size_t n_rows = rows_.size();
+ if (!n_rows) return NULLPTR;
+
+ // Build the arrays column-by-column from the rows
+ std::vector<std::shared_ptr<arrow::Array>> arrays(output_schema->num_fields());
+ for (size_t i_table = 0; i_table < n_tables_; ++i_table) {
+ int n_src_cols = state.at(i_table)->get_schema()->num_fields();
+ {
+ for (col_index_t i_src_col = 0; i_src_col < n_src_cols; ++i_src_col) {
+ util::optional<col_index_t> i_dst_col_opt =
+ state[i_table]->map_src_to_dst(i_src_col);
+ if (!i_dst_col_opt) continue;
+ col_index_t i_dst_col = *i_dst_col_opt;
+ const auto& src_field = state[i_table]->get_schema()->field(i_src_col);
+ const auto& dst_field = output_schema->field(i_dst_col);
+ assert(src_field->type()->Equals(dst_field->type()));
+ assert(src_field->name() == dst_field->name());
+ const auto& field_type = src_field->type();
+
+ if (field_type->Equals(arrow::int32())) {
+ ARROW_ASSIGN_OR_RAISE(
+ arrays.at(i_dst_col),
+ (materialize_primitive_column<arrow::Int32Builder, int32_t>(i_table,
+ i_src_col)));
+ } else if (field_type->Equals(arrow::int64())) {
+ ARROW_ASSIGN_OR_RAISE(
+ arrays.at(i_dst_col),
+ (materialize_primitive_column<arrow::Int64Builder, int64_t>(i_table,
+ i_src_col)));
+ } else if (field_type->Equals(arrow::float64())) {
+ ARROW_ASSIGN_OR_RAISE(
+ arrays.at(i_dst_col),
+ (materialize_primitive_column<arrow::DoubleBuilder, double>(i_table,
+ i_src_col)));
+ } else {
+ ARROW_RETURN_NOT_OK(
+ Status::Invalid("Unsupported data type: ", src_field->name()));
+ }
+ }
+ }
+ }
+
+ // Build the result
+ assert(sizeof(size_t) >= sizeof(int64_t)); // Make takes signed int64_t for num_rows
+
+ // TODO: check n_rows for cast
+ std::shared_ptr<arrow::RecordBatch> r =
+ arrow::RecordBatch::Make(output_schema, (int64_t)n_rows, arrays);
+ return r;
+ }
+
+ // Returns true if there are no rows
+ bool empty() const { return rows_.empty(); }
+
+ private:
+ // Contains shared_ptr refs for all RecordBatches referred to by the contents of rows_
+ std::unordered_map<uintptr_t, std::shared_ptr<RecordBatch>> _ptr2ref;
+
+ // Row table references
+ std::vector<CompositeReferenceRow<MAX_TABLES>> rows_;
+
+ // Total number of tables in the composite table
+ size_t n_tables_;
+
+ // Adds a RecordBatch ref to the mapping, if needed
+ void add_record_batch_ref(const std::shared_ptr<RecordBatch>& ref) {
+ if (!_ptr2ref.count((uintptr_t)ref.get())) _ptr2ref[(uintptr_t)ref.get()] = ref;
+ }
+
+ template <class Builder, class PrimitiveType>
+ Result<std::shared_ptr<Array>> materialize_primitive_column(size_t i_table,
+ col_index_t i_col) {
+ Builder builder;
+ ARROW_RETURN_NOT_OK(builder.Reserve(rows_.size()));
+ for (row_index_t i_row = 0; i_row < rows_.size(); ++i_row) {
+ const auto& ref = rows_[i_row].refs[i_table];
+ if (ref.batch) {
+ builder.UnsafeAppend(
+ ref.batch->column_data(i_col)->template GetValues<PrimitiveType>(1)[ref.row]);
+ } else {
+ builder.UnsafeAppendNull();
+ }
+ }
+ std::shared_ptr<Array> result;
+ ARROW_RETURN_NOT_OK(builder.Finish(&result));
+ return result;
+ }
+};
+
+class AsofJoinNode : public ExecNode {
+ // Constructs labels for inputs
+ static std::vector<std::string> build_input_labels(
+ const std::vector<ExecNode*>& inputs) {
+ std::vector<std::string> r(inputs.size());
+ for (size_t i = 0; i < r.size(); ++i) r[i] = "input_" + std::to_string(i) + "_label";
+ return r;
+ }
+
+ // Advances the RHS as far as possible to be up to date for the current LHS timestamp
+ bool update_rhs() {
+ auto& lhs = *_state.at(0);
+ auto lhs_latest_time = lhs.get_latest_time();
+ bool any_updated = false;
+ for (size_t i = 1; i < _state.size(); ++i)
+ any_updated |= _state[i]->advance_and_memoize(lhs_latest_time);
+ return any_updated;
+ }
+
+ // Returns false if RHS not up to date for LHS
+ bool is_up_to_date_for_lhs_row() const {
+ auto& lhs = *_state[0];
+ if (lhs.empty()) return false; // can't proceed if nothing on the LHS
+ int64_t lhs_ts = lhs.get_latest_time();
+ for (size_t i = 1; i < _state.size(); ++i) {
+ auto& rhs = *_state[i];
+ if (!rhs.finished()) {
+ // If RHS is finished, then we know it's up to date (but if it isn't, it might be
+ // up to date)
+ if (rhs.empty())
+ return false; // RHS isn't finished, but is empty --> not up to date
+ if (lhs_ts >= rhs.get_latest_time())
+ return false; // TS not up to date (and not finished)
+ }
+ }
+ return true;
+ }
+
+ Result<std::shared_ptr<RecordBatch>> process_inner() {
+ assert(!_state.empty());
+ auto& lhs = *_state.at(0);
+
+ // Construct new target table if needed
+ CompositeReferenceTable<MAX_JOIN_TABLES> dst(_state.size());
+
+ // 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;
+
+ // Advance each of the RHS as far as possible to be up to date for the LHS timestamp
+ bool any_advanced = update_rhs();
+
+ // Only update if we have up-to-date information for the LHS row
+ if (is_up_to_date_for_lhs_row()) {
+ dst.emplace(_state, _options.tolerance);
+ if (!lhs.advance()) break; // if we can't advance LHS, we're done for this batch
+ } else {
+ if ((!any_advanced) && (_state.size() > 1)) break; // need to wait for new data
+ }
+ }
+
+ // Prune memo entries that have expired (to bound memory consumption)
+ if (!lhs.empty()) {
+ for (size_t i = 1; i < _state.size(); ++i) {
+ _state[i]->remove_memo_entries_with_lesser_time(lhs.get_latest_time() -
+ _options.tolerance);
+ }
+ }
+
+ // Emit the batch
+ if (dst.empty()) {
+ return NULLPTR;
+ } else {
+ return dst.materialize(output_schema(), _state);
+ }
+ }
+
+ void process() {
+ std::cerr << "process() begin\n";
+
+ std::lock_guard<std::mutex> guard(_gate);
+ if (finished_.is_finished()) {
+ std::cerr << "InputReceived EARLYEND\n";
+ return;
+ }
+
+ // Process batches while we have data
+ for (;;) {
+ Result<std::shared_ptr<RecordBatch>> result = process_inner();
+
+ if (result.ok()) {
+ auto out_rb = *result;
+ if (!out_rb) break;
+ ++_progress_batches_produced;
+ ExecBatch out_b(*out_rb);
+ outputs_[0]->InputReceived(this, std::move(out_b));
+ } else {
+ StopProducing();
+ ErrorIfNotOk(result.status());
+ return;
+ }
+ }
+
+ std::cerr << "process() end\n";
+
+ // 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 (_state.at(0)->finished()) {
+ total_batches_produced_ = util::make_optional<int>(_progress_batches_produced);
+ StopProducing();
+ assert(total_batches_produced_.has_value());
Review Comment:
```suggestion
DCHECK(total_batches_produced_.has_value());
```
##########
cpp/src/arrow/compute/exec/asof_join_node.cc:
##########
@@ -0,0 +1,800 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+#include <iostream>
+#include <unordered_map>
+
+#include <arrow/api.h>
+#include <arrow/compute/api.h>
+#include <arrow/util/optional.h>
+#include "arrow/compute/exec/asof_join.h"
+#include "arrow/compute/exec/exec_plan.h"
+#include "arrow/compute/exec/options.h"
+#include "arrow/compute/exec/schema_util.h"
+#include "arrow/compute/exec/util.h"
+#include "arrow/result.h"
+#include "arrow/status.h"
+#include "arrow/util/checked_cast.h"
+#include "arrow/util/future.h"
+#include "arrow/util/make_unique.h"
+
+#include <condition_variable>
+#include <mutex>
+#include <thread>
+
+namespace arrow {
+namespace compute {
+
+/**
+ * Simple implementation for an unbound concurrent queue
+ */
+template <class T>
+class ConcurrentQueue {
+ public:
+ T pop() {
+ std::unique_lock<std::mutex> lock(mutex_);
+ cond_.wait(lock, [&] { return !queue_.empty(); });
+ auto item = queue_.front();
+ queue_.pop();
+ return item;
+ }
+
+ void push(const T& item) {
+ std::unique_lock<std::mutex> lock(mutex_);
+ queue_.push(item);
+ cond_.notify_one();
+ }
+
+ util::optional<T> try_pop() {
+ // Try to pop the oldest value from the queue (or return nullopt if none)
+ std::unique_lock<std::mutex> lock(mutex_);
+ if (queue_.empty()) {
+ return util::nullopt;
+ } else {
+ auto item = queue_.front();
+ queue_.pop();
+ return item;
+ }
+ }
+
+ bool empty() const {
+ std::unique_lock<std::mutex> lock(mutex_);
+ return queue_.empty();
+ }
+
+ // Un-synchronized access to front
+ // For this to be "safe":
+ // 1) the caller logically guarantees that queue is not empty
+ // 2) pop/try_pop cannot be called concurrently with this
+ const T& unsync_front() const { return queue_.front(); }
+
+ private:
+ std::queue<T> queue_;
+ mutable std::mutex mutex_;
+ std::condition_variable cond_;
+};
+
+struct MemoStore {
+ // Stores last known values for all the keys
+
+ struct Entry {
+ // Timestamp associated with the entry
+ int64_t _time;
+
+ // Batch associated with the entry (perf is probably OK for this; batches change
+ // rarely)
+ std::shared_ptr<arrow::RecordBatch> _batch;
+
+ // Row associated with the entry
+ row_index_t _row;
+ };
+
+ std::unordered_map<KeyType, Entry> _entries;
+
+ void store(const std::shared_ptr<RecordBatch>& batch, row_index_t row, int64_t time,
+ KeyType key) {
+ auto& e = _entries[key];
+ // that we can do this assignment optionally, is why we
+ // can get array 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;
+ }
+
+ util::optional<const Entry*> get_entry_for_key(KeyType key) const {
+ auto e = _entries.find(key);
+ if (_entries.end() == e) return util::nullopt;
+ return util::optional<const Entry*>(&e->second);
+ }
+
+ void remove_entries_with_lesser_time(int64_t ts) {
+ size_t dbg_size0 = _entries.size();
+ for (auto e = _entries.begin(); e != _entries.end();)
+ if (e->second._time < ts)
+ e = _entries.erase(e);
+ else
+ ++e;
+ size_t dbg_size1 = _entries.size();
+ if (dbg_size1 < dbg_size0) {
+ // cerr << "Removed " << dbg_size0-dbg_size1 << " memo entries.\n";
+ }
+ }
+};
+
+class InputState {
+ // InputState correponds to an input
+ // Input record batches are queued up in InputState until processed and
+ // turned into output record batches.
+
+ public:
+ InputState(const std::shared_ptr<arrow::Schema>& schema,
+ const std::string& time_col_name, const std::string& key_col_name,
+ util::optional<KeyType> wildcard_key)
+ : queue_(),
+ wildcard_key_(wildcard_key),
+ schema_(schema),
+ time_col_index_(
+ schema->GetFieldIndex(time_col_name)), // TODO: handle missing field name
+ key_col_index_(schema->GetFieldIndex(key_col_name)) {}
+
+ col_index_t init_src_to_dst_mapping(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 && is_time_or_key_column(i)))
+ src_to_dst_[i] = dst_offset++;
+ return dst_offset;
+ }
+
+ const util::optional<col_index_t>& map_src_to_dst(col_index_t src) const {
+ return src_to_dst_[src];
+ }
+
+ bool is_time_or_key_column(col_index_t i) const {
+ assert(i < schema_->num_fields());
+ return (i == time_col_index_) || (i == key_col_index_);
+ }
+
+ // Gets the latest row index, assuming the queue isn't empty
+ row_index_t get_latest_row() const { return latest_ref_row_; }
+
+ bool empty() const {
+ if (latest_ref_row_ > 0)
+ return false; // cannot be empty if ref row is >0 -- can avoid slow queue lock
+ // below
+ return queue_.empty();
+ }
+
+ int countbatches_processed_() const { return batches_processed_; }
+ int count_total_batches() const { return total_batches_; }
+
+ // Gets latest batch (precondition: must not be empty)
+ const std::shared_ptr<arrow::RecordBatch>& get_latest_batch() const {
+ return queue_.unsync_front();
+ }
+ KeyType get_latest_key() const {
+ return queue_.unsync_front()
+ ->column_data(key_col_index_)
+ ->GetValues<KeyType>(1)[latest_ref_row_];
+ }
+ int64_t get_latest_time() const {
+ return queue_.unsync_front()
+ ->column_data(time_col_index_)
+ ->GetValues<int64_t>(1)[latest_ref_row_];
+ }
+
+ bool finished() const { return batches_processed_ == total_batches_; }
+
+ bool advance() {
+ // 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();
+ if (have_active_batch) {
+ // If we have an active batch
+ if (++latest_ref_row_ >= (row_index_t)queue_.unsync_front()->num_rows()) {
+ // hit the end of the batch, need to get the next batch if possible.
+ ++batches_processed_;
+ latest_ref_row_ = 0;
+ have_active_batch &= !queue_.try_pop();
+ if (have_active_batch)
+ assert(queue_.unsync_front()->num_rows() > 0); // empty batches disallowed
+ }
+ }
+ return have_active_batch;
+ }
+
+ // Advance the data to be immediately past the specified TS, updating latest and
+ // latest_ref_row to the latest data prior to that immediate just past Returns true if
+ // updates were made, false if not.
+ bool advance_and_memoize(int64_t ts) {
+ // Advance the right side row index until we reach the latest right row (for each key)
+ // for the given left timestamp.
+
+ // Check if already updated for TS (or if there is no latest)
+ if (empty()) return false; // can't advance if empty
+ auto latest_time = get_latest_time();
+ if (latest_time > ts) return false; // already advanced
+
+ // Not updated. Try to update and possibly advance.
+ bool updated = false;
+ do {
+ latest_time = get_latest_time();
+ // 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 <= ts) {
+ memo_.store(get_latest_batch(), latest_ref_row_, latest_time, get_latest_key());
+ } else {
+ break; // hit a future timestamp -- done updating for now
+ }
+ updated = true;
+ } while (advance());
+ return updated;
+ }
+
+ void push(const std::shared_ptr<arrow::RecordBatch>& rb) {
+ if (rb->num_rows() > 0) {
+ queue_.push(rb);
+ } else {
+ ++batches_processed_; // don't enqueue empty batches, just record as processed
+ }
+ }
+
+ util::optional<const MemoStore::Entry*> get_memo_entry_for_key(KeyType key) {
+ auto r = memo_.get_entry_for_key(key);
+ if (r.has_value()) return r;
+ if (wildcard_key_.has_value()) r = memo_.get_entry_for_key(*wildcard_key_);
+ return r;
+ }
+
+ util::optional<int64_t> get_memo_time_for_key(KeyType key) {
+ auto r = get_memo_entry_for_key(key);
+ return r.has_value() ? util::make_optional((*r)->_time) : util::nullopt;
+ }
+
+ void remove_memo_entries_with_lesser_time(int64_t ts) {
+ memo_.remove_entries_with_lesser_time(ts);
+ }
+
+ const std::shared_ptr<Schema>& get_schema() const { return schema_; }
+
+ void set_total_batches(int n) {
+ assert(n >= 0);
+ assert(total_batches_ == -1); // shouldn't be set more than once
+ total_batches_ = n;
+ }
+
+ private:
+ // Pending record batches. The latest is the front. Batches cannot be empty.
+ ConcurrentQueue<std::shared_ptr<RecordBatch>> queue_;
+
+ // Wildcard key for this input, if applicable.
+ util::optional<KeyType> wildcard_key_;
+
+ // Schema associated with the input
+ std::shared_ptr<Schema> schema_;
+
+ // Total number of batches (only int because InputFinished uses int)
+ int total_batches_ = -1;
+
+ // Number of batches processed so far (only int because InputFinished uses int)
+ int batches_processed_ = 0;
+
+ // Index of the time col
+ col_index_t time_col_index_;
+
+ // Index of the key col
+ col_index_t key_col_index_;
+
+ // Index of the latest row reference within; if >0 then queue_ cannot be empty
+ row_index_t latest_ref_row_ =
+ 0; // must be < queue_.front()->num_rows() if queue_ is non-empty
+
+ // Stores latest known values for the various keys
+ MemoStore memo_;
+
+ // Mapping of source columns to destination columns
+ std::vector<util::optional<col_index_t>> src_to_dst_;
+};
+
+template <size_t MAX_TABLES>
+struct CompositeReferenceRow {
+ struct Entry {
+ arrow::RecordBatch* batch; // can be NULL if there's no value
+ row_index_t row;
+ };
+ Entry refs[MAX_TABLES];
+};
+
+// A table of composite reference rows. Rows maintain pointers to the
+// constituent record batches, but the overall table retains shared_ptr
+// references to ensure memory remains resident while the table is live.
+//
+// The main reason for this is that, especially for wide tables, joins
+// are effectively row-oriented, rather than column-oriented. Separating
+// the join part from the columnar materialization part simplifies the
+// logic around data types and increases efficiency.
+//
+// We don't put the shared_ptr's into the rows for efficiency reasons.
+template <size_t MAX_TABLES>
+class CompositeReferenceTable {
+ public:
+ explicit CompositeReferenceTable(size_t n_tables) : n_tables_(n_tables) {
+ assert(n_tables_ >= 1);
+ assert(n_tables_ <= MAX_TABLES);
+ }
+
+ size_t n_rows() const { return rows_.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<std::unique_ptr<InputState>>& in, int64_t tolerance) {
+ assert(in.size() == n_tables_);
+
+ // Get the LHS key
+ KeyType key = in[0]->get_latest_key();
+
+ // Add row and setup LHS
+ // (the LHS state comes just from the latest row of the LHS table)
+ assert(!in[0]->empty());
+ const std::shared_ptr<arrow::RecordBatch>& lhs_latest_batch =
+ in[0]->get_latest_batch();
+ row_index_t lhs_latest_row = in[0]->get_latest_row();
+ int64_t lhs_latest_time = in[0]->get_latest_time();
+ 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 = rows_.size() + new_batch_size;
+ if (rows_.capacity() < new_capacity) rows_.reserve(new_capacity);
+ }
+ rows_.resize(rows_.size() + 1);
+ auto& row = rows_.back();
+ row.refs[0].batch = lhs_latest_batch.get();
+ row.refs[0].row = lhs_latest_row;
+ add_record_batch_ref(lhs_latest_batch);
+
+ // 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) {
+ util::optional<const MemoStore::Entry*> opt_entry =
+ in[i]->get_memo_entry_for_key(key);
+ if (opt_entry.has_value()) {
+ assert(*opt_entry);
+ if ((*opt_entry)->_time + tolerance >= lhs_latest_time) {
+ // Have a valid entry
+ const MemoStore::Entry* entry = *opt_entry;
+ row.refs[i].batch = entry->_batch.get();
+ row.refs[i].row = entry->_row;
+ add_record_batch_ref(entry->_batch);
+ continue;
+ }
+ }
+ row.refs[i].batch = NULL;
+ row.refs[i].row = 0;
+ }
+ }
+
+ // Materializes the current reference table into a target record batch
+ Result<std::shared_ptr<RecordBatch>> materialize(
+ const std::shared_ptr<arrow::Schema>& output_schema,
+ const std::vector<std::unique_ptr<InputState>>& state) {
+ // cerr << "materialize BEGIN\n";
+ assert(state.size() == n_tables_);
+ assert(state.size() >= 1);
+
+ // Don't build empty batches
+ size_t n_rows = rows_.size();
+ if (!n_rows) return NULLPTR;
+
+ // Build the arrays column-by-column from the rows
+ std::vector<std::shared_ptr<arrow::Array>> arrays(output_schema->num_fields());
+ for (size_t i_table = 0; i_table < n_tables_; ++i_table) {
+ int n_src_cols = state.at(i_table)->get_schema()->num_fields();
+ {
+ for (col_index_t i_src_col = 0; i_src_col < n_src_cols; ++i_src_col) {
+ util::optional<col_index_t> i_dst_col_opt =
+ state[i_table]->map_src_to_dst(i_src_col);
+ if (!i_dst_col_opt) continue;
+ col_index_t i_dst_col = *i_dst_col_opt;
+ const auto& src_field = state[i_table]->get_schema()->field(i_src_col);
+ const auto& dst_field = output_schema->field(i_dst_col);
+ assert(src_field->type()->Equals(dst_field->type()));
+ assert(src_field->name() == dst_field->name());
+ const auto& field_type = src_field->type();
+
+ if (field_type->Equals(arrow::int32())) {
+ ARROW_ASSIGN_OR_RAISE(
+ arrays.at(i_dst_col),
+ (materialize_primitive_column<arrow::Int32Builder, int32_t>(i_table,
+ i_src_col)));
+ } else if (field_type->Equals(arrow::int64())) {
+ ARROW_ASSIGN_OR_RAISE(
+ arrays.at(i_dst_col),
+ (materialize_primitive_column<arrow::Int64Builder, int64_t>(i_table,
+ i_src_col)));
+ } else if (field_type->Equals(arrow::float64())) {
+ ARROW_ASSIGN_OR_RAISE(
+ arrays.at(i_dst_col),
+ (materialize_primitive_column<arrow::DoubleBuilder, double>(i_table,
+ i_src_col)));
+ } else {
+ ARROW_RETURN_NOT_OK(
+ Status::Invalid("Unsupported data type: ", src_field->name()));
+ }
+ }
+ }
+ }
+
+ // Build the result
+ assert(sizeof(size_t) >= sizeof(int64_t)); // Make takes signed int64_t for num_rows
+
+ // TODO: check n_rows for cast
+ std::shared_ptr<arrow::RecordBatch> r =
+ arrow::RecordBatch::Make(output_schema, (int64_t)n_rows, arrays);
+ return r;
+ }
+
+ // Returns true if there are no rows
+ bool empty() const { return rows_.empty(); }
+
+ private:
+ // Contains shared_ptr refs for all RecordBatches referred to by the contents of rows_
+ std::unordered_map<uintptr_t, std::shared_ptr<RecordBatch>> _ptr2ref;
+
+ // Row table references
+ std::vector<CompositeReferenceRow<MAX_TABLES>> rows_;
+
+ // Total number of tables in the composite table
+ size_t n_tables_;
+
+ // Adds a RecordBatch ref to the mapping, if needed
+ void add_record_batch_ref(const std::shared_ptr<RecordBatch>& ref) {
+ if (!_ptr2ref.count((uintptr_t)ref.get())) _ptr2ref[(uintptr_t)ref.get()] = ref;
+ }
+
+ template <class Builder, class PrimitiveType>
+ Result<std::shared_ptr<Array>> materialize_primitive_column(size_t i_table,
+ col_index_t i_col) {
+ Builder builder;
+ ARROW_RETURN_NOT_OK(builder.Reserve(rows_.size()));
+ for (row_index_t i_row = 0; i_row < rows_.size(); ++i_row) {
+ const auto& ref = rows_[i_row].refs[i_table];
+ if (ref.batch) {
+ builder.UnsafeAppend(
+ ref.batch->column_data(i_col)->template GetValues<PrimitiveType>(1)[ref.row]);
+ } else {
+ builder.UnsafeAppendNull();
+ }
+ }
+ std::shared_ptr<Array> result;
+ ARROW_RETURN_NOT_OK(builder.Finish(&result));
+ return result;
+ }
+};
+
+class AsofJoinNode : public ExecNode {
+ // Constructs labels for inputs
+ static std::vector<std::string> build_input_labels(
+ const std::vector<ExecNode*>& inputs) {
+ std::vector<std::string> r(inputs.size());
+ for (size_t i = 0; i < r.size(); ++i) r[i] = "input_" + std::to_string(i) + "_label";
+ return r;
+ }
+
+ // Advances the RHS as far as possible to be up to date for the current LHS timestamp
+ bool update_rhs() {
+ auto& lhs = *_state.at(0);
+ auto lhs_latest_time = lhs.get_latest_time();
+ bool any_updated = false;
+ for (size_t i = 1; i < _state.size(); ++i)
+ any_updated |= _state[i]->advance_and_memoize(lhs_latest_time);
+ return any_updated;
+ }
+
+ // Returns false if RHS not up to date for LHS
+ bool is_up_to_date_for_lhs_row() const {
+ auto& lhs = *_state[0];
+ if (lhs.empty()) return false; // can't proceed if nothing on the LHS
+ int64_t lhs_ts = lhs.get_latest_time();
+ for (size_t i = 1; i < _state.size(); ++i) {
+ auto& rhs = *_state[i];
+ if (!rhs.finished()) {
+ // If RHS is finished, then we know it's up to date (but if it isn't, it might be
+ // up to date)
+ if (rhs.empty())
+ return false; // RHS isn't finished, but is empty --> not up to date
+ if (lhs_ts >= rhs.get_latest_time())
+ return false; // TS not up to date (and not finished)
+ }
+ }
+ return true;
+ }
+
+ Result<std::shared_ptr<RecordBatch>> process_inner() {
+ assert(!_state.empty());
+ auto& lhs = *_state.at(0);
+
+ // Construct new target table if needed
+ CompositeReferenceTable<MAX_JOIN_TABLES> dst(_state.size());
+
+ // 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;
+
+ // Advance each of the RHS as far as possible to be up to date for the LHS timestamp
+ bool any_advanced = update_rhs();
+
+ // Only update if we have up-to-date information for the LHS row
+ if (is_up_to_date_for_lhs_row()) {
+ dst.emplace(_state, _options.tolerance);
+ if (!lhs.advance()) break; // if we can't advance LHS, we're done for this batch
+ } else {
+ if ((!any_advanced) && (_state.size() > 1)) break; // need to wait for new data
+ }
+ }
+
+ // Prune memo entries that have expired (to bound memory consumption)
+ if (!lhs.empty()) {
+ for (size_t i = 1; i < _state.size(); ++i) {
+ _state[i]->remove_memo_entries_with_lesser_time(lhs.get_latest_time() -
+ _options.tolerance);
+ }
+ }
+
+ // Emit the batch
+ if (dst.empty()) {
+ return NULLPTR;
+ } else {
+ return dst.materialize(output_schema(), _state);
+ }
+ }
+
+ void process() {
+ std::cerr << "process() begin\n";
+
+ std::lock_guard<std::mutex> guard(_gate);
+ if (finished_.is_finished()) {
+ std::cerr << "InputReceived EARLYEND\n";
+ return;
+ }
+
+ // Process batches while we have data
+ for (;;) {
+ Result<std::shared_ptr<RecordBatch>> result = process_inner();
+
+ if (result.ok()) {
+ auto out_rb = *result;
+ if (!out_rb) break;
+ ++_progress_batches_produced;
+ ExecBatch out_b(*out_rb);
+ outputs_[0]->InputReceived(this, std::move(out_b));
+ } else {
+ StopProducing();
+ ErrorIfNotOk(result.status());
+ return;
+ }
+ }
+
+ std::cerr << "process() end\n";
+
+ // 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 (_state.at(0)->finished()) {
+ total_batches_produced_ = util::make_optional<int>(_progress_batches_produced);
+ StopProducing();
+ assert(total_batches_produced_.has_value());
+ outputs_[0]->InputFinished(this, *total_batches_produced_);
+ }
+ }
+
+ void process_thread() {
+ std::cerr << "AsofJoinNode::process_thread started.\n";
+ for (;;) {
+ if (!_process.pop()) {
+ std::cerr << "AsofJoinNode::process_thread done.\n";
+ return;
+ }
+ process();
+ }
+ }
+
+ static void process_thread_wrapper(AsofJoinNode* node) { node->process_thread(); }
+
+ public:
+ AsofJoinNode(ExecPlan* plan, NodeVector inputs, std::vector<std::string> input_labels,
+ const AsofJoinNodeOptions& join_options,
+ std::shared_ptr<Schema> output_schema,
+ std::unique_ptr<AsofJoinSchema> schema_mgr);
+
+ virtual ~AsofJoinNode() {
+ _process.push(false); // poison pill
+ _process_thread.join();
+ }
+
+ static arrow::Result<ExecNode*> Make(ExecPlan* plan, std::vector<ExecNode*> inputs,
+ const ExecNodeOptions& options) {
+ std::unique_ptr<AsofJoinSchema> schema_mgr =
+ ::arrow::internal::make_unique<AsofJoinSchema>();
+
+ const auto& join_options = checked_cast<const AsofJoinNodeOptions&>(options);
+ std::shared_ptr<Schema> output_schema =
+ schema_mgr->MakeOutputSchema(inputs, join_options);
+
+ std::vector<std::string> input_labels(inputs.size());
+ input_labels[0] = "left";
+ for (size_t i = 1; i < inputs.size(); ++i) {
+ input_labels[i] = "right_" + std::to_string(i);
+ }
+
+ return plan->EmplaceNode<AsofJoinNode>(plan, inputs, std::move(input_labels),
+ join_options, std::move(output_schema),
+ std::move(schema_mgr));
+ }
+
+ const char* kind_name() const override { return "AsofJoinNode"; }
+
+ void InputReceived(ExecNode* input, ExecBatch batch) override {
+ // Get the input
+ ARROW_DCHECK(std::find(inputs_.begin(), inputs_.end(), input) != inputs_.end());
+ size_t k = std::find(inputs_.begin(), inputs_.end(), input) - inputs_.begin();
+ std::cerr << "InputReceived BEGIN (k=" << k << ")\n";
+
+ // Put into the queue
+ auto rb = *batch.ToRecordBatch(input->output_schema());
+
+ _state.at(k)->push(rb);
+ _process.push(true);
Review Comment:
Doesn't this need to be guarded with a mutex?
##########
cpp/src/arrow/compute/exec/asof_join_node.cc:
##########
@@ -0,0 +1,800 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+#include <iostream>
+#include <unordered_map>
+
+#include <arrow/api.h>
+#include <arrow/compute/api.h>
+#include <arrow/util/optional.h>
+#include "arrow/compute/exec/asof_join.h"
+#include "arrow/compute/exec/exec_plan.h"
+#include "arrow/compute/exec/options.h"
+#include "arrow/compute/exec/schema_util.h"
+#include "arrow/compute/exec/util.h"
+#include "arrow/result.h"
+#include "arrow/status.h"
+#include "arrow/util/checked_cast.h"
+#include "arrow/util/future.h"
+#include "arrow/util/make_unique.h"
+
+#include <condition_variable>
+#include <mutex>
+#include <thread>
+
+namespace arrow {
+namespace compute {
+
+/**
+ * Simple implementation for an unbound concurrent queue
+ */
+template <class T>
+class ConcurrentQueue {
+ public:
+ T pop() {
+ std::unique_lock<std::mutex> lock(mutex_);
+ cond_.wait(lock, [&] { return !queue_.empty(); });
+ auto item = queue_.front();
+ queue_.pop();
+ return item;
+ }
+
+ void push(const T& item) {
+ std::unique_lock<std::mutex> lock(mutex_);
+ queue_.push(item);
+ cond_.notify_one();
+ }
+
+ util::optional<T> try_pop() {
+ // Try to pop the oldest value from the queue (or return nullopt if none)
+ std::unique_lock<std::mutex> lock(mutex_);
+ if (queue_.empty()) {
+ return util::nullopt;
+ } else {
+ auto item = queue_.front();
+ queue_.pop();
+ return item;
+ }
+ }
+
+ bool empty() const {
+ std::unique_lock<std::mutex> lock(mutex_);
+ return queue_.empty();
+ }
+
+ // Un-synchronized access to front
+ // For this to be "safe":
+ // 1) the caller logically guarantees that queue is not empty
+ // 2) pop/try_pop cannot be called concurrently with this
+ const T& unsync_front() const { return queue_.front(); }
+
+ private:
+ std::queue<T> queue_;
+ mutable std::mutex mutex_;
+ std::condition_variable cond_;
+};
+
+struct MemoStore {
+ // Stores last known values for all the keys
+
+ struct Entry {
+ // Timestamp associated with the entry
+ int64_t _time;
+
+ // Batch associated with the entry (perf is probably OK for this; batches change
+ // rarely)
+ std::shared_ptr<arrow::RecordBatch> _batch;
+
+ // Row associated with the entry
+ row_index_t _row;
+ };
+
+ std::unordered_map<KeyType, Entry> _entries;
+
+ void store(const std::shared_ptr<RecordBatch>& batch, row_index_t row, int64_t time,
+ KeyType key) {
+ auto& e = _entries[key];
+ // that we can do this assignment optionally, is why we
+ // can get array 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;
+ }
+
+ util::optional<const Entry*> get_entry_for_key(KeyType key) const {
+ auto e = _entries.find(key);
+ if (_entries.end() == e) return util::nullopt;
+ return util::optional<const Entry*>(&e->second);
+ }
+
+ void remove_entries_with_lesser_time(int64_t ts) {
+ size_t dbg_size0 = _entries.size();
+ for (auto e = _entries.begin(); e != _entries.end();)
+ if (e->second._time < ts)
+ e = _entries.erase(e);
+ else
+ ++e;
+ size_t dbg_size1 = _entries.size();
+ if (dbg_size1 < dbg_size0) {
+ // cerr << "Removed " << dbg_size0-dbg_size1 << " memo entries.\n";
+ }
+ }
+};
+
+class InputState {
+ // InputState correponds to an input
+ // Input record batches are queued up in InputState until processed and
+ // turned into output record batches.
+
+ public:
+ InputState(const std::shared_ptr<arrow::Schema>& schema,
+ const std::string& time_col_name, const std::string& key_col_name,
+ util::optional<KeyType> wildcard_key)
+ : queue_(),
+ wildcard_key_(wildcard_key),
+ schema_(schema),
+ time_col_index_(
+ schema->GetFieldIndex(time_col_name)), // TODO: handle missing field name
+ key_col_index_(schema->GetFieldIndex(key_col_name)) {}
+
+ col_index_t init_src_to_dst_mapping(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 && is_time_or_key_column(i)))
+ src_to_dst_[i] = dst_offset++;
+ return dst_offset;
+ }
+
+ const util::optional<col_index_t>& map_src_to_dst(col_index_t src) const {
+ return src_to_dst_[src];
+ }
+
+ bool is_time_or_key_column(col_index_t i) const {
+ assert(i < schema_->num_fields());
+ return (i == time_col_index_) || (i == key_col_index_);
+ }
+
+ // Gets the latest row index, assuming the queue isn't empty
+ row_index_t get_latest_row() const { return latest_ref_row_; }
+
+ bool empty() const {
+ if (latest_ref_row_ > 0)
+ return false; // cannot be empty if ref row is >0 -- can avoid slow queue lock
+ // below
+ return queue_.empty();
+ }
+
+ int countbatches_processed_() const { return batches_processed_; }
+ int count_total_batches() const { return total_batches_; }
+
+ // Gets latest batch (precondition: must not be empty)
+ const std::shared_ptr<arrow::RecordBatch>& get_latest_batch() const {
+ return queue_.unsync_front();
+ }
+ KeyType get_latest_key() const {
+ return queue_.unsync_front()
+ ->column_data(key_col_index_)
+ ->GetValues<KeyType>(1)[latest_ref_row_];
+ }
+ int64_t get_latest_time() const {
+ return queue_.unsync_front()
+ ->column_data(time_col_index_)
+ ->GetValues<int64_t>(1)[latest_ref_row_];
+ }
+
+ bool finished() const { return batches_processed_ == total_batches_; }
+
+ bool advance() {
+ // 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();
+ if (have_active_batch) {
+ // If we have an active batch
+ if (++latest_ref_row_ >= (row_index_t)queue_.unsync_front()->num_rows()) {
+ // hit the end of the batch, need to get the next batch if possible.
+ ++batches_processed_;
+ latest_ref_row_ = 0;
+ have_active_batch &= !queue_.try_pop();
+ if (have_active_batch)
+ assert(queue_.unsync_front()->num_rows() > 0); // empty batches disallowed
+ }
+ }
+ return have_active_batch;
+ }
+
+ // Advance the data to be immediately past the specified TS, updating latest and
+ // latest_ref_row to the latest data prior to that immediate just past Returns true if
+ // updates were made, false if not.
+ bool advance_and_memoize(int64_t ts) {
+ // Advance the right side row index until we reach the latest right row (for each key)
+ // for the given left timestamp.
+
+ // Check if already updated for TS (or if there is no latest)
+ if (empty()) return false; // can't advance if empty
+ auto latest_time = get_latest_time();
+ if (latest_time > ts) return false; // already advanced
+
+ // Not updated. Try to update and possibly advance.
+ bool updated = false;
+ do {
+ latest_time = get_latest_time();
+ // 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 <= ts) {
+ memo_.store(get_latest_batch(), latest_ref_row_, latest_time, get_latest_key());
+ } else {
+ break; // hit a future timestamp -- done updating for now
+ }
+ updated = true;
+ } while (advance());
+ return updated;
+ }
+
+ void push(const std::shared_ptr<arrow::RecordBatch>& rb) {
+ if (rb->num_rows() > 0) {
+ queue_.push(rb);
+ } else {
+ ++batches_processed_; // don't enqueue empty batches, just record as processed
+ }
+ }
+
+ util::optional<const MemoStore::Entry*> get_memo_entry_for_key(KeyType key) {
+ auto r = memo_.get_entry_for_key(key);
+ if (r.has_value()) return r;
+ if (wildcard_key_.has_value()) r = memo_.get_entry_for_key(*wildcard_key_);
+ return r;
+ }
+
+ util::optional<int64_t> get_memo_time_for_key(KeyType key) {
+ auto r = get_memo_entry_for_key(key);
+ return r.has_value() ? util::make_optional((*r)->_time) : util::nullopt;
+ }
+
+ void remove_memo_entries_with_lesser_time(int64_t ts) {
+ memo_.remove_entries_with_lesser_time(ts);
+ }
+
+ const std::shared_ptr<Schema>& get_schema() const { return schema_; }
+
+ void set_total_batches(int n) {
+ assert(n >= 0);
+ assert(total_batches_ == -1); // shouldn't be set more than once
+ total_batches_ = n;
+ }
+
+ private:
+ // Pending record batches. The latest is the front. Batches cannot be empty.
+ ConcurrentQueue<std::shared_ptr<RecordBatch>> queue_;
+
+ // Wildcard key for this input, if applicable.
+ util::optional<KeyType> wildcard_key_;
+
+ // Schema associated with the input
+ std::shared_ptr<Schema> schema_;
+
+ // Total number of batches (only int because InputFinished uses int)
+ int total_batches_ = -1;
+
+ // Number of batches processed so far (only int because InputFinished uses int)
+ int batches_processed_ = 0;
+
+ // Index of the time col
+ col_index_t time_col_index_;
+
+ // Index of the key col
+ col_index_t key_col_index_;
+
+ // Index of the latest row reference within; if >0 then queue_ cannot be empty
+ row_index_t latest_ref_row_ =
+ 0; // must be < queue_.front()->num_rows() if queue_ is non-empty
+
+ // Stores latest known values for the various keys
+ MemoStore memo_;
+
+ // Mapping of source columns to destination columns
+ std::vector<util::optional<col_index_t>> src_to_dst_;
+};
+
+template <size_t MAX_TABLES>
+struct CompositeReferenceRow {
+ struct Entry {
+ arrow::RecordBatch* batch; // can be NULL if there's no value
+ row_index_t row;
+ };
+ Entry refs[MAX_TABLES];
+};
+
+// A table of composite reference rows. Rows maintain pointers to the
+// constituent record batches, but the overall table retains shared_ptr
+// references to ensure memory remains resident while the table is live.
+//
+// The main reason for this is that, especially for wide tables, joins
+// are effectively row-oriented, rather than column-oriented. Separating
+// the join part from the columnar materialization part simplifies the
+// logic around data types and increases efficiency.
+//
+// We don't put the shared_ptr's into the rows for efficiency reasons.
+template <size_t MAX_TABLES>
+class CompositeReferenceTable {
+ public:
+ explicit CompositeReferenceTable(size_t n_tables) : n_tables_(n_tables) {
+ assert(n_tables_ >= 1);
+ assert(n_tables_ <= MAX_TABLES);
+ }
+
+ size_t n_rows() const { return rows_.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<std::unique_ptr<InputState>>& in, int64_t tolerance) {
+ assert(in.size() == n_tables_);
+
+ // Get the LHS key
+ KeyType key = in[0]->get_latest_key();
+
+ // Add row and setup LHS
+ // (the LHS state comes just from the latest row of the LHS table)
+ assert(!in[0]->empty());
+ const std::shared_ptr<arrow::RecordBatch>& lhs_latest_batch =
+ in[0]->get_latest_batch();
+ row_index_t lhs_latest_row = in[0]->get_latest_row();
+ int64_t lhs_latest_time = in[0]->get_latest_time();
+ 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 = rows_.size() + new_batch_size;
+ if (rows_.capacity() < new_capacity) rows_.reserve(new_capacity);
+ }
+ rows_.resize(rows_.size() + 1);
+ auto& row = rows_.back();
+ row.refs[0].batch = lhs_latest_batch.get();
+ row.refs[0].row = lhs_latest_row;
+ add_record_batch_ref(lhs_latest_batch);
+
+ // 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) {
+ util::optional<const MemoStore::Entry*> opt_entry =
+ in[i]->get_memo_entry_for_key(key);
+ if (opt_entry.has_value()) {
+ assert(*opt_entry);
+ if ((*opt_entry)->_time + tolerance >= lhs_latest_time) {
+ // Have a valid entry
+ const MemoStore::Entry* entry = *opt_entry;
+ row.refs[i].batch = entry->_batch.get();
+ row.refs[i].row = entry->_row;
+ add_record_batch_ref(entry->_batch);
+ continue;
+ }
+ }
+ row.refs[i].batch = NULL;
+ row.refs[i].row = 0;
+ }
+ }
+
+ // Materializes the current reference table into a target record batch
+ Result<std::shared_ptr<RecordBatch>> materialize(
+ const std::shared_ptr<arrow::Schema>& output_schema,
+ const std::vector<std::unique_ptr<InputState>>& state) {
+ // cerr << "materialize BEGIN\n";
+ assert(state.size() == n_tables_);
+ assert(state.size() >= 1);
+
+ // Don't build empty batches
+ size_t n_rows = rows_.size();
+ if (!n_rows) return NULLPTR;
+
+ // Build the arrays column-by-column from the rows
+ std::vector<std::shared_ptr<arrow::Array>> arrays(output_schema->num_fields());
+ for (size_t i_table = 0; i_table < n_tables_; ++i_table) {
+ int n_src_cols = state.at(i_table)->get_schema()->num_fields();
+ {
+ for (col_index_t i_src_col = 0; i_src_col < n_src_cols; ++i_src_col) {
+ util::optional<col_index_t> i_dst_col_opt =
+ state[i_table]->map_src_to_dst(i_src_col);
+ if (!i_dst_col_opt) continue;
+ col_index_t i_dst_col = *i_dst_col_opt;
+ const auto& src_field = state[i_table]->get_schema()->field(i_src_col);
+ const auto& dst_field = output_schema->field(i_dst_col);
+ assert(src_field->type()->Equals(dst_field->type()));
+ assert(src_field->name() == dst_field->name());
+ const auto& field_type = src_field->type();
+
+ if (field_type->Equals(arrow::int32())) {
+ ARROW_ASSIGN_OR_RAISE(
+ arrays.at(i_dst_col),
+ (materialize_primitive_column<arrow::Int32Builder, int32_t>(i_table,
+ i_src_col)));
+ } else if (field_type->Equals(arrow::int64())) {
+ ARROW_ASSIGN_OR_RAISE(
+ arrays.at(i_dst_col),
+ (materialize_primitive_column<arrow::Int64Builder, int64_t>(i_table,
+ i_src_col)));
+ } else if (field_type->Equals(arrow::float64())) {
+ ARROW_ASSIGN_OR_RAISE(
+ arrays.at(i_dst_col),
+ (materialize_primitive_column<arrow::DoubleBuilder, double>(i_table,
+ i_src_col)));
+ } else {
+ ARROW_RETURN_NOT_OK(
+ Status::Invalid("Unsupported data type: ", src_field->name()));
+ }
+ }
+ }
+ }
+
+ // Build the result
+ assert(sizeof(size_t) >= sizeof(int64_t)); // Make takes signed int64_t for num_rows
+
+ // TODO: check n_rows for cast
+ std::shared_ptr<arrow::RecordBatch> r =
+ arrow::RecordBatch::Make(output_schema, (int64_t)n_rows, arrays);
+ return r;
+ }
+
+ // Returns true if there are no rows
+ bool empty() const { return rows_.empty(); }
+
+ private:
+ // Contains shared_ptr refs for all RecordBatches referred to by the contents of rows_
+ std::unordered_map<uintptr_t, std::shared_ptr<RecordBatch>> _ptr2ref;
+
+ // Row table references
+ std::vector<CompositeReferenceRow<MAX_TABLES>> rows_;
+
+ // Total number of tables in the composite table
+ size_t n_tables_;
+
+ // Adds a RecordBatch ref to the mapping, if needed
+ void add_record_batch_ref(const std::shared_ptr<RecordBatch>& ref) {
+ if (!_ptr2ref.count((uintptr_t)ref.get())) _ptr2ref[(uintptr_t)ref.get()] = ref;
+ }
+
+ template <class Builder, class PrimitiveType>
+ Result<std::shared_ptr<Array>> materialize_primitive_column(size_t i_table,
+ col_index_t i_col) {
+ Builder builder;
+ ARROW_RETURN_NOT_OK(builder.Reserve(rows_.size()));
+ for (row_index_t i_row = 0; i_row < rows_.size(); ++i_row) {
+ const auto& ref = rows_[i_row].refs[i_table];
+ if (ref.batch) {
+ builder.UnsafeAppend(
+ ref.batch->column_data(i_col)->template GetValues<PrimitiveType>(1)[ref.row]);
+ } else {
+ builder.UnsafeAppendNull();
+ }
+ }
+ std::shared_ptr<Array> result;
+ ARROW_RETURN_NOT_OK(builder.Finish(&result));
+ return result;
+ }
+};
+
+class AsofJoinNode : public ExecNode {
+ // Constructs labels for inputs
+ static std::vector<std::string> build_input_labels(
+ const std::vector<ExecNode*>& inputs) {
+ std::vector<std::string> r(inputs.size());
+ for (size_t i = 0; i < r.size(); ++i) r[i] = "input_" + std::to_string(i) + "_label";
+ return r;
+ }
+
+ // Advances the RHS as far as possible to be up to date for the current LHS timestamp
+ bool update_rhs() {
+ auto& lhs = *_state.at(0);
+ auto lhs_latest_time = lhs.get_latest_time();
+ bool any_updated = false;
+ for (size_t i = 1; i < _state.size(); ++i)
+ any_updated |= _state[i]->advance_and_memoize(lhs_latest_time);
+ return any_updated;
+ }
+
+ // Returns false if RHS not up to date for LHS
+ bool is_up_to_date_for_lhs_row() const {
+ auto& lhs = *_state[0];
+ if (lhs.empty()) return false; // can't proceed if nothing on the LHS
+ int64_t lhs_ts = lhs.get_latest_time();
+ for (size_t i = 1; i < _state.size(); ++i) {
+ auto& rhs = *_state[i];
+ if (!rhs.finished()) {
+ // If RHS is finished, then we know it's up to date (but if it isn't, it might be
+ // up to date)
+ if (rhs.empty())
+ return false; // RHS isn't finished, but is empty --> not up to date
+ if (lhs_ts >= rhs.get_latest_time())
+ return false; // TS not up to date (and not finished)
+ }
+ }
+ return true;
+ }
+
+ Result<std::shared_ptr<RecordBatch>> process_inner() {
+ assert(!_state.empty());
+ auto& lhs = *_state.at(0);
+
+ // Construct new target table if needed
+ CompositeReferenceTable<MAX_JOIN_TABLES> dst(_state.size());
+
+ // 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;
+
+ // Advance each of the RHS as far as possible to be up to date for the LHS timestamp
+ bool any_advanced = update_rhs();
+
+ // Only update if we have up-to-date information for the LHS row
+ if (is_up_to_date_for_lhs_row()) {
+ dst.emplace(_state, _options.tolerance);
+ if (!lhs.advance()) break; // if we can't advance LHS, we're done for this batch
+ } else {
+ if ((!any_advanced) && (_state.size() > 1)) break; // need to wait for new data
+ }
+ }
+
+ // Prune memo entries that have expired (to bound memory consumption)
+ if (!lhs.empty()) {
+ for (size_t i = 1; i < _state.size(); ++i) {
+ _state[i]->remove_memo_entries_with_lesser_time(lhs.get_latest_time() -
+ _options.tolerance);
+ }
+ }
+
+ // Emit the batch
+ if (dst.empty()) {
+ return NULLPTR;
+ } else {
+ return dst.materialize(output_schema(), _state);
+ }
+ }
+
+ void process() {
+ std::cerr << "process() begin\n";
+
+ std::lock_guard<std::mutex> guard(_gate);
+ if (finished_.is_finished()) {
+ std::cerr << "InputReceived EARLYEND\n";
+ return;
+ }
+
+ // Process batches while we have data
+ for (;;) {
+ Result<std::shared_ptr<RecordBatch>> result = process_inner();
+
+ if (result.ok()) {
+ auto out_rb = *result;
+ if (!out_rb) break;
+ ++_progress_batches_produced;
+ ExecBatch out_b(*out_rb);
+ outputs_[0]->InputReceived(this, std::move(out_b));
+ } else {
+ StopProducing();
+ ErrorIfNotOk(result.status());
+ return;
+ }
+ }
+
+ std::cerr << "process() end\n";
+
+ // 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 (_state.at(0)->finished()) {
+ total_batches_produced_ = util::make_optional<int>(_progress_batches_produced);
Review Comment:
Does this need to be stored in a member variable?
##########
cpp/src/arrow/compute/exec/asof_join_node.cc:
##########
@@ -0,0 +1,800 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+#include <iostream>
+#include <unordered_map>
+
+#include <arrow/api.h>
+#include <arrow/compute/api.h>
+#include <arrow/util/optional.h>
+#include "arrow/compute/exec/asof_join.h"
+#include "arrow/compute/exec/exec_plan.h"
+#include "arrow/compute/exec/options.h"
+#include "arrow/compute/exec/schema_util.h"
+#include "arrow/compute/exec/util.h"
+#include "arrow/result.h"
+#include "arrow/status.h"
+#include "arrow/util/checked_cast.h"
+#include "arrow/util/future.h"
+#include "arrow/util/make_unique.h"
+
+#include <condition_variable>
+#include <mutex>
+#include <thread>
+
+namespace arrow {
+namespace compute {
+
+/**
+ * Simple implementation for an unbound concurrent queue
+ */
+template <class T>
+class ConcurrentQueue {
+ public:
+ T pop() {
+ std::unique_lock<std::mutex> lock(mutex_);
+ cond_.wait(lock, [&] { return !queue_.empty(); });
+ auto item = queue_.front();
+ queue_.pop();
+ return item;
+ }
+
+ void push(const T& item) {
+ std::unique_lock<std::mutex> lock(mutex_);
+ queue_.push(item);
+ cond_.notify_one();
+ }
+
+ util::optional<T> try_pop() {
+ // Try to pop the oldest value from the queue (or return nullopt if none)
+ std::unique_lock<std::mutex> lock(mutex_);
+ if (queue_.empty()) {
+ return util::nullopt;
+ } else {
+ auto item = queue_.front();
+ queue_.pop();
+ return item;
+ }
+ }
+
+ bool empty() const {
+ std::unique_lock<std::mutex> lock(mutex_);
+ return queue_.empty();
+ }
+
+ // Un-synchronized access to front
+ // For this to be "safe":
+ // 1) the caller logically guarantees that queue is not empty
+ // 2) pop/try_pop cannot be called concurrently with this
+ const T& unsync_front() const { return queue_.front(); }
+
+ private:
+ std::queue<T> queue_;
+ mutable std::mutex mutex_;
+ std::condition_variable cond_;
+};
+
+struct MemoStore {
+ // Stores last known values for all the keys
+
+ struct Entry {
+ // Timestamp associated with the entry
+ int64_t _time;
+
+ // Batch associated with the entry (perf is probably OK for this; batches change
+ // rarely)
+ std::shared_ptr<arrow::RecordBatch> _batch;
+
+ // Row associated with the entry
+ row_index_t _row;
+ };
+
+ std::unordered_map<KeyType, Entry> _entries;
+
+ void store(const std::shared_ptr<RecordBatch>& batch, row_index_t row, int64_t time,
+ KeyType key) {
+ auto& e = _entries[key];
+ // that we can do this assignment optionally, is why we
+ // can get array 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;
+ }
+
+ util::optional<const Entry*> get_entry_for_key(KeyType key) const {
+ auto e = _entries.find(key);
+ if (_entries.end() == e) return util::nullopt;
+ return util::optional<const Entry*>(&e->second);
+ }
+
+ void remove_entries_with_lesser_time(int64_t ts) {
+ size_t dbg_size0 = _entries.size();
+ for (auto e = _entries.begin(); e != _entries.end();)
+ if (e->second._time < ts)
+ e = _entries.erase(e);
+ else
+ ++e;
+ size_t dbg_size1 = _entries.size();
+ if (dbg_size1 < dbg_size0) {
+ // cerr << "Removed " << dbg_size0-dbg_size1 << " memo entries.\n";
+ }
+ }
+};
+
+class InputState {
+ // InputState correponds to an input
+ // Input record batches are queued up in InputState until processed and
+ // turned into output record batches.
+
+ public:
+ InputState(const std::shared_ptr<arrow::Schema>& schema,
+ const std::string& time_col_name, const std::string& key_col_name,
+ util::optional<KeyType> wildcard_key)
+ : queue_(),
+ wildcard_key_(wildcard_key),
+ schema_(schema),
+ time_col_index_(
+ schema->GetFieldIndex(time_col_name)), // TODO: handle missing field name
+ key_col_index_(schema->GetFieldIndex(key_col_name)) {}
+
+ col_index_t init_src_to_dst_mapping(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 && is_time_or_key_column(i)))
+ src_to_dst_[i] = dst_offset++;
+ return dst_offset;
+ }
+
+ const util::optional<col_index_t>& map_src_to_dst(col_index_t src) const {
+ return src_to_dst_[src];
+ }
+
+ bool is_time_or_key_column(col_index_t i) const {
+ assert(i < schema_->num_fields());
+ return (i == time_col_index_) || (i == key_col_index_);
+ }
+
+ // Gets the latest row index, assuming the queue isn't empty
+ row_index_t get_latest_row() const { return latest_ref_row_; }
+
+ bool empty() const {
+ if (latest_ref_row_ > 0)
+ return false; // cannot be empty if ref row is >0 -- can avoid slow queue lock
+ // below
+ return queue_.empty();
+ }
+
+ int countbatches_processed_() const { return batches_processed_; }
+ int count_total_batches() const { return total_batches_; }
+
+ // Gets latest batch (precondition: must not be empty)
+ const std::shared_ptr<arrow::RecordBatch>& get_latest_batch() const {
+ return queue_.unsync_front();
+ }
+ KeyType get_latest_key() const {
+ return queue_.unsync_front()
+ ->column_data(key_col_index_)
+ ->GetValues<KeyType>(1)[latest_ref_row_];
+ }
+ int64_t get_latest_time() const {
+ return queue_.unsync_front()
+ ->column_data(time_col_index_)
+ ->GetValues<int64_t>(1)[latest_ref_row_];
+ }
+
+ bool finished() const { return batches_processed_ == total_batches_; }
+
+ bool advance() {
+ // 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();
+ if (have_active_batch) {
+ // If we have an active batch
+ if (++latest_ref_row_ >= (row_index_t)queue_.unsync_front()->num_rows()) {
+ // hit the end of the batch, need to get the next batch if possible.
+ ++batches_processed_;
+ latest_ref_row_ = 0;
+ have_active_batch &= !queue_.try_pop();
+ if (have_active_batch)
+ assert(queue_.unsync_front()->num_rows() > 0); // empty batches disallowed
+ }
+ }
+ return have_active_batch;
+ }
+
+ // Advance the data to be immediately past the specified TS, updating latest and
+ // latest_ref_row to the latest data prior to that immediate just past Returns true if
+ // updates were made, false if not.
+ bool advance_and_memoize(int64_t ts) {
+ // Advance the right side row index until we reach the latest right row (for each key)
+ // for the given left timestamp.
+
+ // Check if already updated for TS (or if there is no latest)
+ if (empty()) return false; // can't advance if empty
+ auto latest_time = get_latest_time();
+ if (latest_time > ts) return false; // already advanced
+
+ // Not updated. Try to update and possibly advance.
+ bool updated = false;
+ do {
+ latest_time = get_latest_time();
+ // 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 <= ts) {
+ memo_.store(get_latest_batch(), latest_ref_row_, latest_time, get_latest_key());
+ } else {
+ break; // hit a future timestamp -- done updating for now
+ }
+ updated = true;
+ } while (advance());
+ return updated;
+ }
+
+ void push(const std::shared_ptr<arrow::RecordBatch>& rb) {
+ if (rb->num_rows() > 0) {
+ queue_.push(rb);
+ } else {
+ ++batches_processed_; // don't enqueue empty batches, just record as processed
+ }
+ }
+
+ util::optional<const MemoStore::Entry*> get_memo_entry_for_key(KeyType key) {
+ auto r = memo_.get_entry_for_key(key);
+ if (r.has_value()) return r;
+ if (wildcard_key_.has_value()) r = memo_.get_entry_for_key(*wildcard_key_);
+ return r;
+ }
+
+ util::optional<int64_t> get_memo_time_for_key(KeyType key) {
+ auto r = get_memo_entry_for_key(key);
+ return r.has_value() ? util::make_optional((*r)->_time) : util::nullopt;
+ }
+
+ void remove_memo_entries_with_lesser_time(int64_t ts) {
+ memo_.remove_entries_with_lesser_time(ts);
+ }
+
+ const std::shared_ptr<Schema>& get_schema() const { return schema_; }
+
+ void set_total_batches(int n) {
+ assert(n >= 0);
+ assert(total_batches_ == -1); // shouldn't be set more than once
+ total_batches_ = n;
+ }
+
+ private:
+ // Pending record batches. The latest is the front. Batches cannot be empty.
+ ConcurrentQueue<std::shared_ptr<RecordBatch>> queue_;
+
+ // Wildcard key for this input, if applicable.
+ util::optional<KeyType> wildcard_key_;
+
+ // Schema associated with the input
+ std::shared_ptr<Schema> schema_;
+
+ // Total number of batches (only int because InputFinished uses int)
+ int total_batches_ = -1;
+
+ // Number of batches processed so far (only int because InputFinished uses int)
+ int batches_processed_ = 0;
+
+ // Index of the time col
+ col_index_t time_col_index_;
+
+ // Index of the key col
+ col_index_t key_col_index_;
+
+ // Index of the latest row reference within; if >0 then queue_ cannot be empty
+ row_index_t latest_ref_row_ =
+ 0; // must be < queue_.front()->num_rows() if queue_ is non-empty
+
+ // Stores latest known values for the various keys
+ MemoStore memo_;
+
+ // Mapping of source columns to destination columns
+ std::vector<util::optional<col_index_t>> src_to_dst_;
+};
+
+template <size_t MAX_TABLES>
+struct CompositeReferenceRow {
+ struct Entry {
+ arrow::RecordBatch* batch; // can be NULL if there's no value
+ row_index_t row;
+ };
+ Entry refs[MAX_TABLES];
+};
+
+// A table of composite reference rows. Rows maintain pointers to the
+// constituent record batches, but the overall table retains shared_ptr
+// references to ensure memory remains resident while the table is live.
+//
+// The main reason for this is that, especially for wide tables, joins
+// are effectively row-oriented, rather than column-oriented. Separating
+// the join part from the columnar materialization part simplifies the
+// logic around data types and increases efficiency.
+//
+// We don't put the shared_ptr's into the rows for efficiency reasons.
+template <size_t MAX_TABLES>
+class CompositeReferenceTable {
+ public:
+ explicit CompositeReferenceTable(size_t n_tables) : n_tables_(n_tables) {
+ assert(n_tables_ >= 1);
+ assert(n_tables_ <= MAX_TABLES);
+ }
+
+ size_t n_rows() const { return rows_.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<std::unique_ptr<InputState>>& in, int64_t tolerance) {
+ assert(in.size() == n_tables_);
+
+ // Get the LHS key
+ KeyType key = in[0]->get_latest_key();
+
+ // Add row and setup LHS
+ // (the LHS state comes just from the latest row of the LHS table)
+ assert(!in[0]->empty());
+ const std::shared_ptr<arrow::RecordBatch>& lhs_latest_batch =
+ in[0]->get_latest_batch();
+ row_index_t lhs_latest_row = in[0]->get_latest_row();
+ int64_t lhs_latest_time = in[0]->get_latest_time();
+ 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 = rows_.size() + new_batch_size;
+ if (rows_.capacity() < new_capacity) rows_.reserve(new_capacity);
+ }
+ rows_.resize(rows_.size() + 1);
+ auto& row = rows_.back();
+ row.refs[0].batch = lhs_latest_batch.get();
+ row.refs[0].row = lhs_latest_row;
+ add_record_batch_ref(lhs_latest_batch);
+
+ // 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) {
+ util::optional<const MemoStore::Entry*> opt_entry =
+ in[i]->get_memo_entry_for_key(key);
+ if (opt_entry.has_value()) {
+ assert(*opt_entry);
+ if ((*opt_entry)->_time + tolerance >= lhs_latest_time) {
+ // Have a valid entry
+ const MemoStore::Entry* entry = *opt_entry;
+ row.refs[i].batch = entry->_batch.get();
+ row.refs[i].row = entry->_row;
+ add_record_batch_ref(entry->_batch);
+ continue;
+ }
+ }
+ row.refs[i].batch = NULL;
+ row.refs[i].row = 0;
+ }
+ }
+
+ // Materializes the current reference table into a target record batch
+ Result<std::shared_ptr<RecordBatch>> materialize(
+ const std::shared_ptr<arrow::Schema>& output_schema,
+ const std::vector<std::unique_ptr<InputState>>& state) {
+ // cerr << "materialize BEGIN\n";
+ assert(state.size() == n_tables_);
+ assert(state.size() >= 1);
+
+ // Don't build empty batches
+ size_t n_rows = rows_.size();
+ if (!n_rows) return NULLPTR;
+
+ // Build the arrays column-by-column from the rows
+ std::vector<std::shared_ptr<arrow::Array>> arrays(output_schema->num_fields());
+ for (size_t i_table = 0; i_table < n_tables_; ++i_table) {
+ int n_src_cols = state.at(i_table)->get_schema()->num_fields();
+ {
+ for (col_index_t i_src_col = 0; i_src_col < n_src_cols; ++i_src_col) {
+ util::optional<col_index_t> i_dst_col_opt =
+ state[i_table]->map_src_to_dst(i_src_col);
+ if (!i_dst_col_opt) continue;
+ col_index_t i_dst_col = *i_dst_col_opt;
+ const auto& src_field = state[i_table]->get_schema()->field(i_src_col);
+ const auto& dst_field = output_schema->field(i_dst_col);
+ assert(src_field->type()->Equals(dst_field->type()));
+ assert(src_field->name() == dst_field->name());
+ const auto& field_type = src_field->type();
+
+ if (field_type->Equals(arrow::int32())) {
+ ARROW_ASSIGN_OR_RAISE(
+ arrays.at(i_dst_col),
+ (materialize_primitive_column<arrow::Int32Builder, int32_t>(i_table,
+ i_src_col)));
+ } else if (field_type->Equals(arrow::int64())) {
+ ARROW_ASSIGN_OR_RAISE(
+ arrays.at(i_dst_col),
+ (materialize_primitive_column<arrow::Int64Builder, int64_t>(i_table,
+ i_src_col)));
+ } else if (field_type->Equals(arrow::float64())) {
+ ARROW_ASSIGN_OR_RAISE(
+ arrays.at(i_dst_col),
+ (materialize_primitive_column<arrow::DoubleBuilder, double>(i_table,
+ i_src_col)));
+ } else {
+ ARROW_RETURN_NOT_OK(
+ Status::Invalid("Unsupported data type: ", src_field->name()));
+ }
+ }
+ }
+ }
+
+ // Build the result
+ assert(sizeof(size_t) >= sizeof(int64_t)); // Make takes signed int64_t for num_rows
+
+ // TODO: check n_rows for cast
+ std::shared_ptr<arrow::RecordBatch> r =
+ arrow::RecordBatch::Make(output_schema, (int64_t)n_rows, arrays);
+ return r;
+ }
+
+ // Returns true if there are no rows
+ bool empty() const { return rows_.empty(); }
+
+ private:
+ // Contains shared_ptr refs for all RecordBatches referred to by the contents of rows_
+ std::unordered_map<uintptr_t, std::shared_ptr<RecordBatch>> _ptr2ref;
+
+ // Row table references
+ std::vector<CompositeReferenceRow<MAX_TABLES>> rows_;
+
+ // Total number of tables in the composite table
+ size_t n_tables_;
+
+ // Adds a RecordBatch ref to the mapping, if needed
+ void add_record_batch_ref(const std::shared_ptr<RecordBatch>& ref) {
+ if (!_ptr2ref.count((uintptr_t)ref.get())) _ptr2ref[(uintptr_t)ref.get()] = ref;
+ }
+
+ template <class Builder, class PrimitiveType>
+ Result<std::shared_ptr<Array>> materialize_primitive_column(size_t i_table,
+ col_index_t i_col) {
+ Builder builder;
+ ARROW_RETURN_NOT_OK(builder.Reserve(rows_.size()));
+ for (row_index_t i_row = 0; i_row < rows_.size(); ++i_row) {
+ const auto& ref = rows_[i_row].refs[i_table];
+ if (ref.batch) {
+ builder.UnsafeAppend(
+ ref.batch->column_data(i_col)->template GetValues<PrimitiveType>(1)[ref.row]);
+ } else {
+ builder.UnsafeAppendNull();
+ }
+ }
+ std::shared_ptr<Array> result;
+ ARROW_RETURN_NOT_OK(builder.Finish(&result));
+ return result;
+ }
+};
+
+class AsofJoinNode : public ExecNode {
+ // Constructs labels for inputs
+ static std::vector<std::string> build_input_labels(
+ const std::vector<ExecNode*>& inputs) {
+ std::vector<std::string> r(inputs.size());
+ for (size_t i = 0; i < r.size(); ++i) r[i] = "input_" + std::to_string(i) + "_label";
+ return r;
+ }
+
+ // Advances the RHS as far as possible to be up to date for the current LHS timestamp
+ bool update_rhs() {
+ auto& lhs = *_state.at(0);
+ auto lhs_latest_time = lhs.get_latest_time();
+ bool any_updated = false;
+ for (size_t i = 1; i < _state.size(); ++i)
+ any_updated |= _state[i]->advance_and_memoize(lhs_latest_time);
+ return any_updated;
+ }
+
+ // Returns false if RHS not up to date for LHS
+ bool is_up_to_date_for_lhs_row() const {
+ auto& lhs = *_state[0];
+ if (lhs.empty()) return false; // can't proceed if nothing on the LHS
+ int64_t lhs_ts = lhs.get_latest_time();
+ for (size_t i = 1; i < _state.size(); ++i) {
+ auto& rhs = *_state[i];
+ if (!rhs.finished()) {
+ // If RHS is finished, then we know it's up to date (but if it isn't, it might be
+ // up to date)
+ if (rhs.empty())
+ return false; // RHS isn't finished, but is empty --> not up to date
+ if (lhs_ts >= rhs.get_latest_time())
+ return false; // TS not up to date (and not finished)
+ }
+ }
+ return true;
+ }
+
+ Result<std::shared_ptr<RecordBatch>> process_inner() {
+ assert(!_state.empty());
+ auto& lhs = *_state.at(0);
+
+ // Construct new target table if needed
+ CompositeReferenceTable<MAX_JOIN_TABLES> dst(_state.size());
+
+ // 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;
+
+ // Advance each of the RHS as far as possible to be up to date for the LHS timestamp
+ bool any_advanced = update_rhs();
+
+ // Only update if we have up-to-date information for the LHS row
+ if (is_up_to_date_for_lhs_row()) {
+ dst.emplace(_state, _options.tolerance);
+ if (!lhs.advance()) break; // if we can't advance LHS, we're done for this batch
+ } else {
+ if ((!any_advanced) && (_state.size() > 1)) break; // need to wait for new data
+ }
+ }
+
+ // Prune memo entries that have expired (to bound memory consumption)
+ if (!lhs.empty()) {
+ for (size_t i = 1; i < _state.size(); ++i) {
+ _state[i]->remove_memo_entries_with_lesser_time(lhs.get_latest_time() -
+ _options.tolerance);
+ }
+ }
+
+ // Emit the batch
+ if (dst.empty()) {
+ return NULLPTR;
+ } else {
+ return dst.materialize(output_schema(), _state);
+ }
+ }
+
+ void process() {
+ std::cerr << "process() begin\n";
+
+ std::lock_guard<std::mutex> guard(_gate);
+ if (finished_.is_finished()) {
+ std::cerr << "InputReceived EARLYEND\n";
+ return;
+ }
+
+ // Process batches while we have data
+ for (;;) {
+ Result<std::shared_ptr<RecordBatch>> result = process_inner();
+
+ if (result.ok()) {
+ auto out_rb = *result;
+ if (!out_rb) break;
+ ++_progress_batches_produced;
+ ExecBatch out_b(*out_rb);
+ outputs_[0]->InputReceived(this, std::move(out_b));
+ } else {
+ StopProducing();
+ ErrorIfNotOk(result.status());
+ return;
+ }
+ }
+
+ std::cerr << "process() end\n";
+
+ // 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 (_state.at(0)->finished()) {
+ total_batches_produced_ = util::make_optional<int>(_progress_batches_produced);
+ StopProducing();
+ assert(total_batches_produced_.has_value());
+ outputs_[0]->InputFinished(this, *total_batches_produced_);
+ }
+ }
+
+ void process_thread() {
+ std::cerr << "AsofJoinNode::process_thread started.\n";
+ for (;;) {
+ if (!_process.pop()) {
+ std::cerr << "AsofJoinNode::process_thread done.\n";
+ return;
+ }
+ process();
+ }
+ }
+
+ static void process_thread_wrapper(AsofJoinNode* node) { node->process_thread(); }
+
+ public:
+ AsofJoinNode(ExecPlan* plan, NodeVector inputs, std::vector<std::string> input_labels,
+ const AsofJoinNodeOptions& join_options,
+ std::shared_ptr<Schema> output_schema,
+ std::unique_ptr<AsofJoinSchema> schema_mgr);
+
+ virtual ~AsofJoinNode() {
+ _process.push(false); // poison pill
+ _process_thread.join();
+ }
+
+ static arrow::Result<ExecNode*> Make(ExecPlan* plan, std::vector<ExecNode*> inputs,
+ const ExecNodeOptions& options) {
+ std::unique_ptr<AsofJoinSchema> schema_mgr =
+ ::arrow::internal::make_unique<AsofJoinSchema>();
+
+ const auto& join_options = checked_cast<const AsofJoinNodeOptions&>(options);
+ std::shared_ptr<Schema> output_schema =
+ schema_mgr->MakeOutputSchema(inputs, join_options);
+
+ std::vector<std::string> input_labels(inputs.size());
+ input_labels[0] = "left";
+ for (size_t i = 1; i < inputs.size(); ++i) {
+ input_labels[i] = "right_" + std::to_string(i);
+ }
+
+ return plan->EmplaceNode<AsofJoinNode>(plan, inputs, std::move(input_labels),
+ join_options, std::move(output_schema),
+ std::move(schema_mgr));
+ }
+
+ const char* kind_name() const override { return "AsofJoinNode"; }
+
+ void InputReceived(ExecNode* input, ExecBatch batch) override {
+ // Get the input
+ ARROW_DCHECK(std::find(inputs_.begin(), inputs_.end(), input) != inputs_.end());
+ size_t k = std::find(inputs_.begin(), inputs_.end(), input) - inputs_.begin();
+ std::cerr << "InputReceived BEGIN (k=" << k << ")\n";
+
+ // Put into the queue
+ auto rb = *batch.ToRecordBatch(input->output_schema());
+
+ _state.at(k)->push(rb);
+ _process.push(true);
+
+ std::cerr << "InputReceived END\n";
+ }
+ void ErrorReceived(ExecNode* input, Status error) override {
+ outputs_[0]->ErrorReceived(this, std::move(error));
+ StopProducing();
+ }
+ void InputFinished(ExecNode* input, int total_batches) override {
+ std::cerr << "InputFinished BEGIN\n";
+ // bool is_finished=false;
+ {
+ std::lock_guard<std::mutex> guard(_gate);
+ std::cerr << "InputFinished find\n";
+ ARROW_DCHECK(std::find(inputs_.begin(), inputs_.end(), input) != inputs_.end());
+ size_t k = std::find(inputs_.begin(), inputs_.end(), input) - inputs_.begin();
+ // cerr << "set_total_batches for input " << k << ": " << total_batches << "\n";
+ _state.at(k)->set_total_batches(total_batches);
+ }
+ // 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.
+ _process.push(true);
+
+ std::cerr << "InputFinished END\n";
+ }
+ Status StartProducing() override {
+ std::cout << "StartProducing"
+ << "\n";
+ finished_ = arrow::Future<>::Make();
+ return Status::OK();
+ }
+ void PauseProducing(ExecNode* output, int32_t counter) override {
+ std::cout << "PauseProducing"
+ << "\n";
+ }
+ void ResumeProducing(ExecNode* output, int32_t counter) override {
+ std::cout << "ResumeProducing"
+ << "\n";
+ }
+ void StopProducing(ExecNode* output) override {
+ DCHECK_EQ(output, outputs_[0]);
+ StopProducing();
+ std::cout << "StopProducing"
+ << "\n";
+ }
+ void StopProducing() override {
+ std::cerr << "StopProducing" << std::endl;
+ // if(batch_count_.Cancel()) finished_.MarkFinished();
+ finished_.MarkFinished();
+ for (auto&& input : inputs_) input->StopProducing(this);
+ }
+ arrow::Future<> finished() override { return finished_; }
+
+ private:
+ std::unique_ptr<AsofJoinSchema> schema_mgr_;
+ arrow::Future<> finished_;
+ // InputStates
+ // Each input state correponds to an input table
+ //
+ std::vector<std::unique_ptr<InputState>> _state;
+ std::mutex _gate;
+ AsofJoinNodeOptions _options;
+
+ // Queue for triggering processing of a given input
+ // (a false value is a poison pill)
+ ConcurrentQueue<bool> _process;
+ // Worker thread
+ std::thread _process_thread;
+
+ // Total batches produced, once we've finished -- only known at completion time.
+ util::optional<int> total_batches_produced_;
+
+ // In-progress batches produced
+ int _progress_batches_produced = 0;
+};
+
+std::shared_ptr<Schema> AsofJoinSchema::MakeOutputSchema(
+ const std::vector<ExecNode*>& inputs, const AsofJoinNodeOptions& options) {
+ std::vector<std::shared_ptr<arrow::Field>> fields;
+ assert(inputs.size() > 1);
Review Comment:
```suggestion
DCHECK_GT(inputs.size(), 1);
```
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org