Skip to content

Commit 9410327

Browse files
authored
GH-44629: [C++][Acero] Use implicit_ordering for asof_join rather than require_sequenced_output (#44616)
### Rationale for this change Changes in #44083 (GH-41706) unnecessarily sequences batches retrieved from scanner where it only requires the batches to provide index according to implicit input order. ### What changes are included in this PR? Setting `implicit_ordering` causes existing code to set batch index, which is then available to the `asof_join` node to sequence the batches int input order. This replaces some of #44083 changes. Some code introduced by #44083 turns out to not be required and has therefore been reverted. ### Are these changes tested? Existing unit tests still pass. ### Are there any user-facing changes? Reverts some breaking user-facing changes done by #44083. * GitHub Issue: #41706 * GitHub Issue: #44629 Authored-by: Enrico Minack <[email protected]> Signed-off-by: Rossi Sun <[email protected]>
1 parent f6e2cbe commit 9410327

File tree

7 files changed

+32
-21
lines changed

7 files changed

+32
-21
lines changed

cpp/src/arrow/acero/options.h

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -103,8 +103,8 @@ class ARROW_ACERO_EXPORT SourceNodeOptions : public ExecNodeOptions {
103103
std::shared_ptr<Schema> output_schema;
104104
/// \brief an asynchronous stream of batches ending with std::nullopt
105105
std::function<Future<std::optional<ExecBatch>>()> generator;
106-
107-
Ordering ordering = Ordering::Unordered();
106+
/// \brief the order of the data, defaults to Ordering::Unordered
107+
Ordering ordering;
108108
};
109109

110110
/// \brief a node that generates data from a table already loaded in memory

cpp/src/arrow/acero/source_node.cc

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -407,7 +407,7 @@ struct SchemaSourceNode : public SourceNode {
407407
struct RecordBatchReaderSourceNode : public SourceNode {
408408
RecordBatchReaderSourceNode(ExecPlan* plan, std::shared_ptr<Schema> schema,
409409
arrow::AsyncGenerator<std::optional<ExecBatch>> generator)
410-
: SourceNode(plan, schema, generator, Ordering::Implicit()) {}
410+
: SourceNode(plan, schema, generator) {}
411411

412412
static Result<ExecNode*> Make(ExecPlan* plan, std::vector<ExecNode*> inputs,
413413
const ExecNodeOptions& options) {

cpp/src/arrow/dataset/scanner.cc

Lines changed: 5 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -1000,6 +1000,7 @@ Result<acero::ExecNode*> MakeScanNode(acero::ExecPlan* plan,
10001000
auto scan_options = scan_node_options.scan_options;
10011001
auto dataset = scan_node_options.dataset;
10021002
bool require_sequenced_output = scan_node_options.require_sequenced_output;
1003+
bool implicit_ordering = scan_node_options.implicit_ordering;
10031004

10041005
RETURN_NOT_OK(NormalizeScanOptions(scan_options, dataset->schema()));
10051006

@@ -1032,11 +1033,11 @@ Result<acero::ExecNode*> MakeScanNode(acero::ExecPlan* plan,
10321033
} else {
10331034
batch_gen = std::move(merged_batch_gen);
10341035
}
1035-
int64_t index = require_sequenced_output ? 0 : compute::kUnsequencedIndex;
1036+
10361037
auto gen = MakeMappedGenerator(
10371038
std::move(batch_gen),
1038-
[scan_options, index](const EnumeratedRecordBatch& partial) mutable
1039-
-> Result<std::optional<compute::ExecBatch>> {
1039+
[scan_options](const EnumeratedRecordBatch& partial)
1040+
-> Result<std::optional<compute::ExecBatch>> {
10401041
// TODO(ARROW-13263) fragments may be able to attach more guarantees to batches
10411042
// than this, for example parquet's row group stats. Failing to do this leaves
10421043
// perf on the table because row group stats could be used to skip kernel execs in
@@ -1057,11 +1058,10 @@ Result<acero::ExecNode*> MakeScanNode(acero::ExecPlan* plan,
10571058
batch->values.emplace_back(partial.record_batch.index);
10581059
batch->values.emplace_back(partial.record_batch.last);
10591060
batch->values.emplace_back(partial.fragment.value->ToString());
1060-
if (index != compute::kUnsequencedIndex) batch->index = index++;
10611061
return batch;
10621062
});
10631063

1064-
auto ordering = require_sequenced_output ? Ordering::Implicit() : Ordering::Unordered();
1064+
auto ordering = implicit_ordering ? Ordering::Implicit() : Ordering::Unordered();
10651065

10661066
auto fields = scan_options->dataset_schema->fields();
10671067
if (scan_options->add_augmented_fields) {

cpp/src/arrow/dataset/scanner.h

Lines changed: 11 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -557,20 +557,27 @@ class ARROW_DS_EXPORT ScannerBuilder {
557557
/// \brief Construct a source ExecNode which yields batches from a dataset scan.
558558
///
559559
/// Does not construct associated filter or project nodes.
560-
/// Yielded batches will be augmented with fragment/batch indices to enable stable
561-
/// ordering for simple ExecPlans.
560+
///
561+
/// Batches are yielded sequentially, like single-threaded,
562+
/// when require_sequenced_output=true.
563+
///
564+
/// Yielded batches will be augmented with fragment/batch indices when
565+
/// implicit_ordering=true to enable stable ordering for simple ExecPlans.
562566
class ARROW_DS_EXPORT ScanNodeOptions : public acero::ExecNodeOptions {
563567
public:
564568
explicit ScanNodeOptions(std::shared_ptr<Dataset> dataset,
565569
std::shared_ptr<ScanOptions> scan_options,
566-
bool require_sequenced_output = false)
570+
bool require_sequenced_output = false,
571+
bool implicit_ordering = false)
567572
: dataset(std::move(dataset)),
568573
scan_options(std::move(scan_options)),
569-
require_sequenced_output(require_sequenced_output) {}
574+
require_sequenced_output(require_sequenced_output),
575+
implicit_ordering(implicit_ordering) {}
570576

571577
std::shared_ptr<Dataset> dataset;
572578
std::shared_ptr<ScanOptions> scan_options;
573579
bool require_sequenced_output;
580+
bool implicit_ordering;
574581
};
575582

576583
/// @}

python/pyarrow/_dataset.pyx

Lines changed: 8 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -4077,13 +4077,15 @@ cdef class _ScanNodeOptions(ExecNodeOptions):
40774077
cdef:
40784078
shared_ptr[CScanOptions] c_scan_options
40794079
bint require_sequenced_output=False
4080+
bint implicit_ordering=False
40804081

40814082
c_scan_options = Scanner._make_scan_options(dataset, scan_options)
40824083

40834084
require_sequenced_output=scan_options.get("require_sequenced_output", False)
4085+
implicit_ordering=scan_options.get("implicit_ordering", False)
40844086

40854087
self.wrapped.reset(
4086-
new CScanNodeOptions(dataset.unwrap(), c_scan_options, require_sequenced_output)
4088+
new CScanNodeOptions(dataset.unwrap(), c_scan_options, require_sequenced_output, implicit_ordering)
40874089
)
40884090

40894091

@@ -4101,8 +4103,8 @@ class ScanNodeOptions(_ScanNodeOptions):
41014103
expression or projection to the scan node that you also supply
41024104
to the filter or project node.
41034105
4104-
Yielded batches will be augmented with fragment/batch indices to
4105-
enable stable ordering for simple ExecPlans.
4106+
Yielded batches will be augmented with fragment/batch indices when
4107+
implicit_ordering=True to enable stable ordering for simple ExecPlans.
41064108
41074109
Parameters
41084110
----------
@@ -4111,7 +4113,9 @@ class ScanNodeOptions(_ScanNodeOptions):
41114113
**kwargs : dict, optional
41124114
Scan options. See `Scanner.from_dataset` for possible arguments.
41134115
require_sequenced_output : bool, default False
4114-
Assert implicit ordering on data.
4116+
Batches are yielded sequentially, like single-threaded
4117+
implicit_ordering : bool, default False
4118+
Preserve implicit ordering of data.
41154119
"""
41164120

41174121
def __init__(self, Dataset dataset, **kwargs):

python/pyarrow/acero.py

Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -56,10 +56,10 @@ class InMemoryDataset:
5656
ds = DatasetModuleStub
5757

5858

59-
def _dataset_to_decl(dataset, use_threads=True, require_sequenced_output=False):
59+
def _dataset_to_decl(dataset, use_threads=True, implicit_ordering=False):
6060
decl = Declaration("scan", ScanNodeOptions(
6161
dataset, use_threads=use_threads,
62-
require_sequenced_output=require_sequenced_output))
62+
implicit_ordering=implicit_ordering))
6363

6464
# Get rid of special dataset columns
6565
# "__fragment_index", "__batch_index", "__last_in_fragment", "__filename"
@@ -316,15 +316,15 @@ def _perform_join_asof(left_operand, left_on, left_by,
316316
left_source = _dataset_to_decl(
317317
left_operand,
318318
use_threads=use_threads,
319-
require_sequenced_output=True)
319+
implicit_ordering=True)
320320
else:
321321
left_source = Declaration(
322322
"table_source", TableSourceNodeOptions(left_operand),
323323
)
324324
if isinstance(right_operand, ds.Dataset):
325325
right_source = _dataset_to_decl(
326326
right_operand, use_threads=use_threads,
327-
require_sequenced_output=True)
327+
implicit_ordering=True)
328328
else:
329329
right_source = Declaration(
330330
"table_source", TableSourceNodeOptions(right_operand)

python/pyarrow/includes/libarrow_dataset.pxd

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -51,7 +51,7 @@ cdef extern from "arrow/dataset/api.h" namespace "arrow::dataset" nogil:
5151
CExpression filter
5252

5353
cdef cppclass CScanNodeOptions "arrow::dataset::ScanNodeOptions"(CExecNodeOptions):
54-
CScanNodeOptions(shared_ptr[CDataset] dataset, shared_ptr[CScanOptions] scan_options, bint require_sequenced_output)
54+
CScanNodeOptions(shared_ptr[CDataset] dataset, shared_ptr[CScanOptions] scan_options, bint require_sequenced_output, bint implicit_ordering)
5555

5656
shared_ptr[CScanOptions] scan_options
5757

0 commit comments

Comments
 (0)