Skip to content

Commit 4f31bfc

Browse files
ManMansonManManson
authored
ARROW-17318: [C++][Dataset] Support async streaming interface for getting fragments in Dataset (apache#13804)
Add `GetFragmentsAsync()` and `GetFragmentsAsyncImpl()` functions to the generic `Dataset` interface, which allows to produce fragments in a streamed fashion. This is one of the prerequisites for making `FileSystemDataset` to support lazy fragment processing, which, in turn, can be used to start scan operations without waiting for the entire dataset to be discovered. To aid the transition process of moving to async implementation in `Dataset`/`AsyncScanner` code, a default implementation for `GetFragmentsAsyncImpl()` is provided (yielding a VectorGenerator over the fragments vector, which is stored by every implementation of Dataset interface at the moment). Tests: unit(release) Signed-off-by: Pavel Solodovnikov <pavel.al.solodovnikov@gmail.com> Authored-by: Pavel Solodovnikov <pavel.al.solodovnikov@gmail.com> Signed-off-by: Antoine Pitrou <antoine@python.org>
1 parent ab71673 commit 4f31bfc

6 files changed

Lines changed: 171 additions & 6 deletions

File tree

‎cpp/src/arrow/dataset/dataset.cc‎

Lines changed: 30 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -15,18 +15,19 @@
1515
// specific language governing permissions and limitations
1616
// under the License.
1717

18-
#include "arrow/dataset/dataset.h"
19-
2018
#include <memory>
2119
#include <utility>
2220

21+
#include "arrow/dataset/dataset.h"
2322
#include "arrow/dataset/dataset_internal.h"
2423
#include "arrow/dataset/scanner.h"
2524
#include "arrow/table.h"
25+
#include "arrow/util/async_generator.h"
2626
#include "arrow/util/bit_util.h"
2727
#include "arrow/util/iterator.h"
2828
#include "arrow/util/logging.h"
2929
#include "arrow/util/make_unique.h"
30+
#include "arrow/util/thread_pool.h"
3031

3132
namespace arrow {
3233

@@ -160,6 +161,33 @@ Result<FragmentIterator> Dataset::GetFragments(compute::Expression predicate) {
160161
: MakeEmptyIterator<std::shared_ptr<Fragment>>();
161162
}
162163

164+
Result<FragmentGenerator> Dataset::GetFragmentsAsync() {
165+
return GetFragmentsAsync(compute::literal(true));
166+
}
167+
168+
Result<FragmentGenerator> Dataset::GetFragmentsAsync(compute::Expression predicate) {
169+
ARROW_ASSIGN_OR_RAISE(
170+
predicate, SimplifyWithGuarantee(std::move(predicate), partition_expression_));
171+
return predicate.IsSatisfiable()
172+
? GetFragmentsAsyncImpl(std::move(predicate),
173+
arrow::internal::GetCpuThreadPool())
174+
: MakeEmptyGenerator<std::shared_ptr<Fragment>>();
175+
}
176+
177+
// Default impl delegating the work to `GetFragmentsImpl` and wrapping it into
178+
// BackgroundGenerator/TransferredGenerator, which offloads potentially
179+
// IO-intensive work to the default IO thread pool and then transfers the control
180+
// back to the specified executor.
181+
Result<FragmentGenerator> Dataset::GetFragmentsAsyncImpl(
182+
compute::Expression predicate, arrow::internal::Executor* executor) {
183+
ARROW_ASSIGN_OR_RAISE(auto iter, GetFragmentsImpl(std::move(predicate)));
184+
ARROW_ASSIGN_OR_RAISE(
185+
auto background_gen,
186+
MakeBackgroundGenerator(std::move(iter), io::default_io_context().executor()));
187+
auto transferred_gen = MakeTransferredGenerator(std::move(background_gen), executor);
188+
return transferred_gen;
189+
}
190+
163191
struct VectorRecordBatchGenerator : InMemoryDataset::RecordBatchGenerator {
164192
explicit VectorRecordBatchGenerator(RecordBatchVector batches)
165193
: batches_(std::move(batches)) {}

‎cpp/src/arrow/dataset/dataset.h‎

Lines changed: 24 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -29,10 +29,16 @@
2929
#include "arrow/compute/exec/expression.h"
3030
#include "arrow/dataset/type_fwd.h"
3131
#include "arrow/dataset/visibility.h"
32+
#include "arrow/util/async_generator_fwd.h"
3233
#include "arrow/util/macros.h"
3334
#include "arrow/util/mutex.h"
3435

3536
namespace arrow {
37+
38+
namespace internal {
39+
class Executor;
40+
} // namespace internal
41+
3642
namespace dataset {
3743

3844
using RecordBatchGenerator = std::function<Future<std::shared_ptr<RecordBatch>>()>;
@@ -134,6 +140,8 @@ class ARROW_DS_EXPORT InMemoryFragment : public Fragment {
134140

135141
/// @}
136142

143+
using FragmentGenerator = AsyncGenerator<std::shared_ptr<Fragment>>;
144+
137145
/// \brief A container of zero or more Fragments.
138146
///
139147
/// A Dataset acts as a union of Fragments, e.g. files deeply nested in a
@@ -148,6 +156,10 @@ class ARROW_DS_EXPORT Dataset : public std::enable_shared_from_this<Dataset> {
148156
Result<FragmentIterator> GetFragments(compute::Expression predicate);
149157
Result<FragmentIterator> GetFragments();
150158

159+
/// \brief Async versions of `GetFragments`.
160+
Result<FragmentGenerator> GetFragmentsAsync(compute::Expression predicate);
161+
Result<FragmentGenerator> GetFragmentsAsync();
162+
151163
const std::shared_ptr<Schema>& schema() const { return schema_; }
152164

153165
/// \brief An expression which evaluates to true for all data viewed by this Dataset.
@@ -174,6 +186,18 @@ class ARROW_DS_EXPORT Dataset : public std::enable_shared_from_this<Dataset> {
174186
Dataset(std::shared_ptr<Schema> schema, compute::Expression partition_expression);
175187

176188
virtual Result<FragmentIterator> GetFragmentsImpl(compute::Expression predicate) = 0;
189+
/// \brief Default non-virtual implementation method for the base
190+
/// `GetFragmentsAsyncImpl` method, which creates a fragment generator for
191+
/// the dataset, possibly filtering results with a predicate (forwarding to
192+
/// the synchronous `GetFragmentsImpl` method and moving the computations
193+
/// to the background, using the IO thread pool).
194+
///
195+
/// Currently, `executor` is always the same as `internal::GetCPUThreadPool()`,
196+
/// which means the results from the underlying fragment generator will be
197+
/// transfered to the default CPU thread pool. The generator itself is
198+
/// offloaded to run on the default IO thread pool.
199+
virtual Result<FragmentGenerator> GetFragmentsAsyncImpl(
200+
compute::Expression predicate, arrow::internal::Executor* executor);
177201

178202
std::shared_ptr<Schema> schema_;
179203
compute::Expression partition_expression_ = compute::literal(true);

‎cpp/src/arrow/dataset/dataset_test.cc‎

Lines changed: 28 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -146,6 +146,34 @@ TEST_F(TestInMemoryDataset, HandlesDifferingSchemas) {
146146
scanner->ToTable());
147147
}
148148

149+
TEST_F(TestInMemoryDataset, GetFragmentsSync) {
150+
constexpr int64_t kBatchSize = 1024;
151+
constexpr int64_t kNumberBatches = 16;
152+
153+
SetSchema({field("i32", int32()), field("f64", float64())});
154+
auto batch = ConstantArrayGenerator::Zeroes(kBatchSize, schema_);
155+
auto reader = ConstantArrayGenerator::Repeat(kNumberBatches, batch);
156+
157+
auto dataset = std::make_shared<InMemoryDataset>(
158+
schema_, RecordBatchVector{static_cast<size_t>(kNumberBatches), batch});
159+
160+
AssertDatasetFragmentsEqual(reader.get(), dataset.get());
161+
}
162+
163+
TEST_F(TestInMemoryDataset, GetFragmentsAsync) {
164+
constexpr int64_t kBatchSize = 1024;
165+
constexpr int64_t kNumberBatches = 16;
166+
167+
SetSchema({field("i32", int32()), field("f64", float64())});
168+
auto batch = ConstantArrayGenerator::Zeroes(kBatchSize, schema_);
169+
auto reader = ConstantArrayGenerator::Repeat(kNumberBatches, batch);
170+
171+
auto dataset = std::make_shared<InMemoryDataset>(
172+
schema_, RecordBatchVector{static_cast<size_t>(kNumberBatches), batch});
173+
174+
AssertDatasetAsyncFragmentsEqual(reader.get(), dataset.get());
175+
}
176+
149177
class TestUnionDataset : public DatasetFixtureMixin {};
150178

151179
TEST_F(TestUnionDataset, ReplaceSchema) {

‎cpp/src/arrow/dataset/test_util.h‎

Lines changed: 17 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -167,7 +167,7 @@ class DatasetFixtureMixin : public ::testing::Test {
167167
void AssertFragmentEquals(RecordBatchReader* expected, Fragment* fragment,
168168
bool ensure_drained = true) {
169169
ASSERT_OK_AND_ASSIGN(auto batch_gen, fragment->ScanBatchesAsync(options_));
170-
AssertScanTaskEquals(expected, batch_gen);
170+
AssertScanTaskEquals(expected, batch_gen, ensure_drained);
171171

172172
if (ensure_drained) {
173173
EnsureRecordBatchReaderDrained(expected);
@@ -191,6 +191,22 @@ class DatasetFixtureMixin : public ::testing::Test {
191191
}
192192
}
193193

194+
void AssertDatasetAsyncFragmentsEqual(RecordBatchReader* expected, Dataset* dataset,
195+
bool ensure_drained = true) {
196+
ASSERT_OK_AND_ASSIGN(auto predicate, options_->filter.Bind(*dataset->schema()));
197+
ASSERT_OK_AND_ASSIGN(auto gen, dataset->GetFragmentsAsync(predicate))
198+
199+
ASSERT_FINISHES_OK(VisitAsyncGenerator(
200+
std::move(gen), [this, expected](const std::shared_ptr<Fragment>& f) {
201+
AssertFragmentEquals(expected, f.get(), false /*ensure_drained*/);
202+
return Status::OK();
203+
}));
204+
205+
if (ensure_drained) {
206+
EnsureRecordBatchReaderDrained(expected);
207+
}
208+
}
209+
194210
/// \brief Ensure that record batches found in reader are equals to the
195211
/// record batches yielded by a scanner.
196212
void AssertScannerEquals(RecordBatchReader* expected, Scanner* scanner,

‎cpp/src/arrow/util/async_generator.h‎

Lines changed: 1 addition & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -25,6 +25,7 @@
2525
#include <optional>
2626
#include <queue>
2727

28+
#include "arrow/util/async_generator_fwd.h"
2829
#include "arrow/util/async_util.h"
2930
#include "arrow/util/functional.h"
3031
#include "arrow/util/future.h"
@@ -66,9 +67,6 @@ namespace arrow {
6667
// until all outstanding futures have completed. Generators that spawn multiple
6768
// concurrent futures may need to hold onto an error while other concurrent futures wrap
6869
// up.
69-
template <typename T>
70-
using AsyncGenerator = std::function<Future<T>()>;
71-
7270
template <typename T>
7371
struct IterationTraits<AsyncGenerator<T>> {
7472
/// \brief by default when iterating through a sequence of AsyncGenerator<T>,

‎cpp/src/arrow/util/async_generator_fwd.h‎

Lines changed: 71 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,71 @@
1+
// Licensed to the Apache Software Foundation (ASF) under one
2+
// or more contributor license agreements. See the NOTICE file
3+
// distributed with this work for additional information
4+
// regarding copyright ownership. The ASF licenses this file
5+
// to you under the Apache License, Version 2.0 (the
6+
// "License"); you may not use this file except in compliance
7+
// with the License. You may obtain a copy of the License at
8+
//
9+
// http://www.apache.org/licenses/LICENSE-2.0
10+
//
11+
// Unless required by applicable law or agreed to in writing,
12+
// software distributed under the License is distributed on an
13+
// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
14+
// KIND, either express or implied. See the License for the
15+
// specific language governing permissions and limitations
16+
// under the License.
17+
18+
#pragma once
19+
20+
#include <functional>
21+
22+
#include "arrow/type_fwd.h"
23+
24+
namespace arrow {
25+
26+
template <typename T>
27+
using AsyncGenerator = std::function<Future<T>()>;
28+
29+
template <typename T, typename V>
30+
class MappingGenerator;
31+
32+
template <typename T, typename ComesAfter, typename IsNext>
33+
class SequencingGenerator;
34+
35+
template <typename T, typename V>
36+
class TransformingGenerator;
37+
38+
template <typename T>
39+
class SerialReadaheadGenerator;
40+
41+
template <typename T>
42+
class ReadaheadGenerator;
43+
44+
template <typename T>
45+
class PushGenerator;
46+
47+
template <typename T>
48+
class MergedGenerator;
49+
50+
template <typename T>
51+
struct Enumerated;
52+
53+
template <typename T>
54+
class EnumeratingGenerator;
55+
56+
template <typename T>
57+
class TransferringGenerator;
58+
59+
template <typename T>
60+
class BackgroundGenerator;
61+
62+
template <typename T>
63+
class GeneratorIterator;
64+
65+
template <typename T>
66+
struct CancellableGenerator;
67+
68+
template <typename T>
69+
class DefaultIfEmptyGenerator;
70+
71+
} // namespace arrow

0 commit comments

Comments
 (0)