Skip to content

Commit f67009a

Browse files
amol-jorisvandenbosschewestonpacepitrou
authored
ARROW-16616: [Python] Add lazy Dataset.filter() method (apache#13409)
Expose a `Dataset.filter` method that applies a filter to the dataset without actually loading it in memory. Addresses what was discussed in apache#13155 (comment) - [x] Update documentation - [x] Ensure the filtered dataset preserves the filter when writing it back - [x] Ensure the filtered dataset preserves the filter when joining - [x] Ensure the filtered dataset preserves the filter when applying standard `Dataset.something` methods. - [x] Allow to extend the filter by adding more conditions subsequently `dataset(filter=X).filter(filter=Y).scanner(filter=Z)` (related to apache#13409 (comment)) - [x] Refactor to use only `Dataset` class instead of `FilteredDataset` as discussed with @ jorisvandenbossche - [x] Add support in replace_schema - [x] Error in get_fragments in case a filter is set. - [x] Verify support in UnionDataset Lead-authored-by: Alessandro Molina <amol@turbogears.org> Co-authored-by: Joris Van den Bossche <jorisvandenbossche@gmail.com> Co-authored-by: Weston Pace <weston.pace@gmail.com> Co-authored-by: Antoine Pitrou <pitrou@free.fr> Signed-off-by: Alessandro Molina <amol@turbogears.org>
1 parent 398ce31 commit f67009a

12 files changed

Lines changed: 339 additions & 44 deletions

File tree

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

Lines changed: 14 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -51,16 +51,27 @@ std::string ToString(JoinType t) {
5151
}
5252

5353
Result<std::shared_ptr<SourceNodeOptions>> SourceNodeOptions::FromTable(
54-
const Table& table, arrow::internal::Executor* exc) {
54+
const Table& table, arrow::internal::Executor* executor) {
5555
std::shared_ptr<RecordBatchReader> reader = std::make_shared<TableBatchReader>(table);
5656

57-
if (exc == nullptr) return Status::TypeError("No executor provided.");
57+
if (executor == nullptr) return Status::TypeError("No executor provided.");
5858

5959
// Map the RecordBatchReader to a SourceNode
60-
ARROW_ASSIGN_OR_RAISE(auto batch_gen, MakeReaderGenerator(std::move(reader), exc));
60+
ARROW_ASSIGN_OR_RAISE(auto batch_gen, MakeReaderGenerator(std::move(reader), executor));
6161

6262
return std::make_shared<SourceNodeOptions>(table.schema(), batch_gen);
6363
}
6464

65+
Result<std::shared_ptr<SourceNodeOptions>> SourceNodeOptions::FromRecordBatchReader(
66+
std::shared_ptr<RecordBatchReader> reader, std::shared_ptr<Schema> schema,
67+
arrow::internal::Executor* executor) {
68+
if (executor == nullptr) return Status::TypeError("No executor provided.");
69+
70+
// Map the RecordBatchReader to a SourceNode
71+
ARROW_ASSIGN_OR_RAISE(auto batch_gen, MakeReaderGenerator(std::move(reader), executor));
72+
73+
return std::make_shared<SourceNodeOptions>(std::move(schema), std::move(batch_gen));
74+
}
75+
6576
} // namespace compute
6677
} // namespace arrow

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

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -65,6 +65,10 @@ class ARROW_EXPORT SourceNodeOptions : public ExecNodeOptions {
6565
static Result<std::shared_ptr<SourceNodeOptions>> FromTable(const Table& table,
6666
arrow::internal::Executor*);
6767

68+
static Result<std::shared_ptr<SourceNodeOptions>> FromRecordBatchReader(
69+
std::shared_ptr<RecordBatchReader> reader, std::shared_ptr<Schema> schema,
70+
arrow::internal::Executor*);
71+
6872
std::shared_ptr<Schema> output_schema;
6973
std::function<Future<std::optional<ExecBatch>>()> generator;
7074
};

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

Lines changed: 7 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -942,12 +942,17 @@ Status ScannerBuilder::Backpressure(compute::BackpressureOptions backpressure) {
942942
return Status::OK();
943943
}
944944

945-
Result<std::shared_ptr<Scanner>> ScannerBuilder::Finish() {
945+
Result<std::shared_ptr<ScanOptions>> ScannerBuilder::GetScanOptions() {
946946
if (!scan_options_->projection.IsBound()) {
947947
RETURN_NOT_OK(Project(scan_options_->dataset_schema->field_names()));
948948
}
949949

950-
return std::make_shared<AsyncScanner>(dataset_, scan_options_);
950+
return scan_options_;
951+
}
952+
953+
Result<std::shared_ptr<Scanner>> ScannerBuilder::Finish() {
954+
ARROW_ASSIGN_OR_RAISE(auto scan_options, GetScanOptions());
955+
return std::make_shared<AsyncScanner>(dataset_, scan_options);
951956
}
952957

953958
namespace {

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

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -534,6 +534,9 @@ class ARROW_DS_EXPORT ScannerBuilder {
534534
/// \brief Override default backpressure configuration
535535
Status Backpressure(compute::BackpressureOptions backpressure);
536536

537+
/// \brief Return the current scan options for the builder.
538+
Result<std::shared_ptr<ScanOptions>> GetScanOptions();
539+
537540
/// \brief Return the constructed now-immutable Scanner object
538541
Result<std::shared_ptr<Scanner>> Finish();
539542

‎docs/source/python/compute.rst‎

Lines changed: 13 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -368,8 +368,19 @@ our ``even_filter`` with a ``pc.field("nums") > 5`` filter:
368368
nums: [[6,8,10]]
369369
chars: [["f","h","l"]]
370370
371-
:class:`.Dataset` currently can be filtered using :meth:`.Dataset.to_table` method
372-
passing a ``filter`` argument. See :ref:`py-filter-dataset` in Dataset documentation.
371+
:class:`.Dataset` can similarly be filtered with the :meth:`.Dataset.filter` method.
372+
The method will return an instance of :class:`.Dataset` which will lazily
373+
apply the filter as soon as actual data of the dataset is accessed:
374+
375+
>>> dataset = ds.dataset(table)
376+
>>> filtered = dataset.filter(pc.field("nums") < 5).filter(pc.field("nums") > 2)
377+
>>> filtered.to_table()
378+
pyarrow.Table
379+
nums: int64
380+
chars: string
381+
----
382+
nums: [[3,4]]
383+
chars: [["c","d"]]
373384

374385

375386
User-Defined Functions

‎python/pyarrow/_dataset.pxd‎

Lines changed: 17 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -47,6 +47,7 @@ cdef class Dataset(_Weakrefable):
4747
cdef:
4848
shared_ptr[CDataset] wrapped
4949
CDataset* dataset
50+
public dict _scan_options
5051

5152
cdef void init(self, const shared_ptr[CDataset]& sp)
5253

@@ -56,6 +57,22 @@ cdef class Dataset(_Weakrefable):
5657
cdef shared_ptr[CDataset] unwrap(self) nogil
5758

5859

60+
cdef class Scanner(_Weakrefable):
61+
cdef:
62+
shared_ptr[CScanner] wrapped
63+
CScanner* scanner
64+
65+
cdef void init(self, const shared_ptr[CScanner]& sp)
66+
67+
@staticmethod
68+
cdef wrap(const shared_ptr[CScanner]& sp)
69+
70+
cdef shared_ptr[CScanner] unwrap(self)
71+
72+
@staticmethod
73+
cdef shared_ptr[CScanOptions] _make_scan_options(Dataset dataset, dict py_scanoptions) except *
74+
75+
5976
cdef class FragmentScanOptions(_Weakrefable):
6077

6178
cdef:

‎python/pyarrow/_dataset.pyx‎

Lines changed: 105 additions & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -156,6 +156,7 @@ cdef class Dataset(_Weakrefable):
156156
cdef void init(self, const shared_ptr[CDataset]& sp):
157157
self.wrapped = sp
158158
self.dataset = sp.get()
159+
self._scan_options = dict()
159160

160161
@staticmethod
161162
cdef wrap(const shared_ptr[CDataset]& sp):
@@ -200,8 +201,14 @@ cdef class Dataset(_Weakrefable):
200201
The new dataset schema.
201202
"""
202203
cdef shared_ptr[CDataset] copy = GetResultValue(
203-
self.dataset.ReplaceSchema(pyarrow_unwrap_schema(schema)))
204-
return Dataset.wrap(move(copy))
204+
self.dataset.ReplaceSchema(pyarrow_unwrap_schema(schema))
205+
)
206+
207+
d = Dataset.wrap(move(copy))
208+
if self._scan_options:
209+
# Preserve scan options if set.
210+
d._scan_options = self._scan_options.copy()
211+
return d
205212

206213
def get_fragments(self, Expression filter=None):
207214
"""Returns an iterator over the fragments in this dataset.
@@ -217,6 +224,18 @@ cdef class Dataset(_Weakrefable):
217224
-------
218225
fragments : iterator of Fragment
219226
"""
227+
if self._scan_options.get("filter") is not None:
228+
# Accessing fragments of a filtered dataset is not supported.
229+
# It would be unclear if you wanted to filter the fragments
230+
# or the rows in those fragments.
231+
raise ValueError(
232+
"Retrieving fragments of a filtered or projected "
233+
"dataset is not allowed. Remove the filtering."
234+
)
235+
236+
return self._get_fragments(filter)
237+
238+
def _get_fragments(self, Expression filter):
220239
cdef:
221240
CExpression c_filter
222241
CFragmentIterator c_iterator
@@ -231,6 +250,24 @@ cdef class Dataset(_Weakrefable):
231250
for maybe_fragment in c_fragments:
232251
yield Fragment.wrap(GetResultValue(move(maybe_fragment)))
233252

253+
def _scanner_options(self, options):
254+
"""Returns the default options to create a new Scanner.
255+
256+
This is automatically invoked by :meth:`Dataset.scanner`
257+
and there is no need to use it.
258+
"""
259+
new_options = options.copy()
260+
261+
# at the moment only support filter
262+
requested_filter = options.get("filter")
263+
current_filter = self._scan_options.get("filter")
264+
if requested_filter is not None and current_filter is not None:
265+
new_options["filter"] = current_filter & requested_filter
266+
elif current_filter is not None:
267+
new_options["filter"] = current_filter
268+
269+
return new_options
270+
234271
def scanner(self, **kwargs):
235272
"""
236273
Build a scan operation against the dataset.
@@ -239,7 +276,7 @@ cdef class Dataset(_Weakrefable):
239276
which exposes further operations (e.g. loading all data as a
240277
table, counting rows).
241278
242-
See the `Scanner.from_dataset` method for further information.
279+
See the :meth:`Scanner.from_dataset` method for further information.
243280
244281
Parameters
245282
----------
@@ -385,6 +422,32 @@ cdef class Dataset(_Weakrefable):
385422
"""The common schema of the full Dataset"""
386423
return pyarrow_wrap_schema(self.dataset.schema())
387424

425+
def filter(self, expression not None):
426+
"""
427+
Apply a row filter to the dataset.
428+
429+
Parameters
430+
----------
431+
expression : Expression
432+
The filter that should be applied to the dataset.
433+
434+
Returns
435+
-------
436+
Dataset
437+
"""
438+
cdef:
439+
Dataset filtered_dataset
440+
441+
new_filter = expression
442+
current_filter = self._scan_options.get("filter")
443+
if current_filter is not None and new_filter is not None:
444+
new_filter = current_filter & new_filter
445+
446+
filtered_dataset = self.__class__.__new__(self.__class__)
447+
filtered_dataset.init(self.wrapped)
448+
filtered_dataset._scan_options = dict(filter=new_filter)
449+
return filtered_dataset
450+
388451
def join(self, right_dataset, keys, right_keys=None, join_type="left outer",
389452
left_suffix=None, right_suffix=None, coalesce_keys=True,
390453
use_threads=True):
@@ -2187,6 +2250,9 @@ cdef class RecordBatchIterator(_Weakrefable):
21872250
self.iterator = make_shared[CRecordBatchIterator](move(iterator))
21882251
return self
21892252

2253+
cdef inline shared_ptr[CRecordBatchIterator] unwrap(self) nogil:
2254+
return self.iterator
2255+
21902256
def __iter__(self):
21912257
return self
21922258

@@ -2354,10 +2420,6 @@ cdef class Scanner(_Weakrefable):
23542420
default pool.
23552421
"""
23562422

2357-
cdef:
2358-
shared_ptr[CScanner] wrapped
2359-
CScanner* scanner
2360-
23612423
def __init__(self):
23622424
_forbid_instantiation(self.__class__)
23632425

@@ -2375,8 +2437,34 @@ cdef class Scanner(_Weakrefable):
23752437
return self.wrapped
23762438

23772439
@staticmethod
2378-
def from_dataset(Dataset dataset not None, *, object columns=None,
2379-
Expression filter=None, int batch_size=_DEFAULT_BATCH_SIZE,
2440+
cdef shared_ptr[CScanOptions] _make_scan_options(Dataset dataset, dict py_scanoptions) except *:
2441+
cdef:
2442+
shared_ptr[CScannerBuilder] builder = make_shared[CScannerBuilder](dataset.unwrap())
2443+
2444+
py_scanoptions = dataset._scanner_options(py_scanoptions)
2445+
2446+
# Need to explicitly expand the arguments as Cython doesn't support
2447+
# keyword expansion in cdef functions.
2448+
_populate_builder(
2449+
builder,
2450+
columns=py_scanoptions.get("columns"),
2451+
filter=py_scanoptions.get("filter"),
2452+
batch_size=py_scanoptions.get("batch_size", _DEFAULT_BATCH_SIZE),
2453+
batch_readahead=py_scanoptions.get(
2454+
"batch_readahead", _DEFAULT_BATCH_READAHEAD),
2455+
fragment_readahead=py_scanoptions.get(
2456+
"fragment_readahead", _DEFAULT_FRAGMENT_READAHEAD),
2457+
use_threads=py_scanoptions.get("use_threads", True),
2458+
memory_pool=py_scanoptions.get("memory_pool"),
2459+
fragment_scan_options=py_scanoptions.get("fragment_scan_options"))
2460+
2461+
return GetResultValue(deref(builder).GetScanOptions())
2462+
2463+
@staticmethod
2464+
def from_dataset(Dataset dataset not None, *,
2465+
object columns=None,
2466+
Expression filter=None,
2467+
int batch_size=_DEFAULT_BATCH_SIZE,
23802468
int batch_readahead=_DEFAULT_BATCH_READAHEAD,
23812469
int fragment_readahead=_DEFAULT_FRAGMENT_READAHEAD,
23822470
FragmentScanOptions fragment_scan_options=None,
@@ -2440,7 +2528,7 @@ cdef class Scanner(_Weakrefable):
24402528
default pool.
24412529
"""
24422530
cdef:
2443-
shared_ptr[CScanOptions] options = make_shared[CScanOptions]()
2531+
shared_ptr[CScanOptions] options
24442532
shared_ptr[CScannerBuilder] builder
24452533
shared_ptr[CScanner] scanner
24462534

@@ -2449,13 +2537,14 @@ cdef class Scanner(_Weakrefable):
24492537
'effect. It will be removed in the next release.',
24502538
FutureWarning)
24512539

2540+
options = Scanner._make_scan_options(
2541+
dataset,
2542+
dict(columns=columns, filter=filter, batch_size=batch_size,
2543+
batch_readahead=batch_readahead,
2544+
fragment_readahead=fragment_readahead, use_threads=use_threads,
2545+
memory_pool=memory_pool, fragment_scan_options=fragment_scan_options)
2546+
)
24522547
builder = make_shared[CScannerBuilder](dataset.unwrap(), options)
2453-
_populate_builder(builder, columns=columns, filter=filter,
2454-
batch_size=batch_size, batch_readahead=batch_readahead,
2455-
fragment_readahead=fragment_readahead, use_threads=use_threads,
2456-
memory_pool=memory_pool,
2457-
fragment_scan_options=fragment_scan_options)
2458-
24592548
scanner = GetResultValue(builder.get().Finish())
24602549
return Scanner.wrap(scanner)
24612550

0 commit comments

Comments
 (0)