Skip to content

Commit 09a794a

Browse files
ARROW-14292: [C++][Python] Join foundation for Tables
This implements the `Table.join` method and the underlying infrastructure. It provides the `tables_join` function that wraps the `execplan` machinery and tests to verify that joins work as expected. Closes apache#12452 from amol-/ARROW-14292 Lead-authored-by: Alessandro Molina <amol@turbogears.org> Co-authored-by: Joris Van den Bossche <jorisvandenbossche@gmail.com> Signed-off-by: Joris Van den Bossche <jorisvandenbossche@gmail.com>
1 parent 4d0436a commit 09a794a

13 files changed

Lines changed: 769 additions & 8 deletions

File tree

‎cpp/src/arrow/CMakeLists.txt‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -395,6 +395,7 @@ if(ARROW_COMPUTE)
395395
compute/exec/options.cc
396396
compute/exec/order_by_impl.cc
397397
compute/exec/partition_util.cc
398+
compute/exec/options.cc
398399
compute/exec/project_node.cc
399400
compute/exec/sink_node.cc
400401
compute/exec/source_node.cc

‎cpp/src/arrow/compute/exec/options.cc‎

Lines changed: 17 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -16,6 +16,10 @@
1616
// under the License.
1717

1818
#include "arrow/compute/exec/options.h"
19+
#include "arrow/compute/exec/exec_plan.h"
20+
#include "arrow/io/util_internal.h"
21+
#include "arrow/table.h"
22+
#include "arrow/util/async_generator.h"
1923
#include "arrow/util/logging.h"
2024

2125
namespace arrow {
@@ -44,5 +48,18 @@ std::string ToString(JoinType t) {
4448
std::abort();
4549
}
4650

51+
Result<std::shared_ptr<SourceNodeOptions>> SourceNodeOptions::FromTable(
52+
const Table& table, arrow::internal::Executor* exc) {
53+
std::shared_ptr<RecordBatchReader> reader = std::make_shared<TableBatchReader>(table);
54+
55+
if (exc == nullptr) return Status::TypeError("No executor provided.");
56+
57+
// Map the RecordBatchReader to a SourceNode
58+
ARROW_ASSIGN_OR_RAISE(auto batch_gen, MakeReaderGenerator(std::move(reader), exc));
59+
60+
return std::shared_ptr<SourceNodeOptions>(
61+
new SourceNodeOptions(table.schema(), batch_gen));
62+
}
63+
4764
} // namespace compute
4865
} // namespace arrow

‎cpp/src/arrow/compute/exec/options.h‎

Lines changed: 7 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -26,16 +26,19 @@
2626
#include "arrow/compute/api_vector.h"
2727
#include "arrow/compute/exec.h"
2828
#include "arrow/compute/exec/expression.h"
29+
#include "arrow/result.h"
30+
#include "arrow/util/async_generator.h"
2931
#include "arrow/util/async_util.h"
3032
#include "arrow/util/optional.h"
3133
#include "arrow/util/visibility.h"
3234

3335
namespace arrow {
3436
namespace compute {
3537

38+
using AsyncExecBatchGenerator = AsyncGenerator<util::optional<ExecBatch>>;
39+
3640
/// \addtogroup execnode-options
3741
/// @{
38-
3942
class ARROW_EXPORT ExecNodeOptions {
4043
public:
4144
virtual ~ExecNodeOptions() = default;
@@ -51,6 +54,9 @@ class ARROW_EXPORT SourceNodeOptions : public ExecNodeOptions {
5154
std::function<Future<util::optional<ExecBatch>>()> generator)
5255
: output_schema(std::move(output_schema)), generator(std::move(generator)) {}
5356

57+
static Result<std::shared_ptr<SourceNodeOptions>> FromTable(const Table& table,
58+
arrow::internal::Executor*);
59+
5460
std::shared_ptr<Schema> output_schema;
5561
std::function<Future<util::optional<ExecBatch>>()> generator;
5662
};

‎python/CMakeLists.txt‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -399,6 +399,7 @@ set(CYTHON_EXTENSIONS
399399
lib
400400
_compute
401401
_csv
402+
_exec_plan
402403
_feather
403404
_fs
404405
_hdfsio

‎python/pyarrow/_compute.pxd‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -49,3 +49,6 @@ cdef class Expression(_Weakrefable):
4949

5050
@staticmethod
5151
cdef Expression _expr_or_scalar(object expr)
52+
53+
54+
cdef CExpression _true

‎python/pyarrow/_compute.pyx‎

Lines changed: 5 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -2218,8 +2218,11 @@ cdef class Expression(_Weakrefable):
22182218
return Expression._call("is_in", [self], options)
22192219

22202220
@staticmethod
2221-
def _field(str name not None):
2222-
return Expression.wrap(CMakeFieldExpression(tobytes(name)))
2221+
def _field(name_or_idx not None):
2222+
if isinstance(name_or_idx, str):
2223+
return Expression.wrap(CMakeFieldExpression(tobytes(name_or_idx)))
2224+
else:
2225+
return Expression.wrap(CMakeFieldExpressionByIndex(name_or_idx))
22232226

22242227
@staticmethod
22252228
def _scalar(value):

0 commit comments

Comments
 (0)