Add row-wise filtering step to read_parquet

[rjzamora]

Closes #13324

Description

This PR adds a post_filters argument to cudf.read_parquet, which is set equal to the filters argument by default. When this argument is set, the specified DNF (disjunctive normal form) filter expression will be applied to the in-memory cudf.DataFrame object after IO is performed.

The overal result of this PR is that the behavior of cudf.read_parquet becomes more consistent with that of pd.read_parquet in the sense that the default result will now enforce filters at a row-wise granularity for both libraries.

Note on the "need" for distinct filters and post_filters arguments

My hope is that post_filters will eventually go away. However, I added a distinct argument for two general reasons:

1. PyArrow does not yet support "is" and "is not" operands in filters. Therefore, we can not pass along all filters from dask/dask_cudf down to cudf.read_parquet using the existing filters argument, because we rely on pyarrow to filter out row-groups (note that dask implements its own filter-conversion utility to avoid this problem). I'm hoping pyarrow will eventually adopt these comparison types (xref: [Python] Add support for "is" and "is not" to pyarrow.parquet.filters_to_expression apache/arrow#34504)
2. When cudf.read_parquet is called from dask_cudf.read_parquet, row-group filtering will have already been applied. Therefore, it is convenient to specify that you only need cudf to provide the post-IO row-wise filtering step. Otherwise, we are effectively duplicating some metadata processing.

My primary concern with adding post_filters is the idea that row-wise filtering could be added at the cuio/libcudf level in the future. In that (hypothetical) case, post_filters wouldn't really be providing any value, but we would probably be able to deprecate it without much pain (if any).

Additional Context

This feature is ultimately needed to support general predicate-pushdown optimizations in Dask Expressions (dask-expr). This is because the high-level optimization logic becomes much simpler when a filter-based operation of a ReadParquet expression can be iteratively "absorbed" into the root ReadParquet expression.

Checklist

• I am familiar with the Contributing Guidelines.
• New or existing tests cover these changes.
• The documentation is up to date with these changes.

[rjzamora]

cc @charlesbluca (for both dask-expr and dask-sql visibility)

[wence-]

I think some of the core logic can be tidied up a bit, and am not sure this works correctly for more than two ANDed or ORed filters.

[wence-]

Is the thought that if the index is not a default rangeindex then it must be intentional and so we must not reset it? How are we to know that the "default" rangeindex is not intentional?

[rjzamora]

I don't have strong feeling about how we handle the default index. I just know that pandas/pyarrow avoids carrying around a filtered index if/when the original DataFrame would have had a default index anyway.

[wence-]

and_ has type bool -> bool -> bool so this only works if you have exactly two conjunctions.

You want (handles length-1 conjunctions as well):

disjunctions.append(functools.reduce(operator.and_, conjunctions))

[rjzamora]

Ah right - I did completely forget that operator.and_/or_ are binops

[wence-]

Again here, wants to be functools.reduce(operator.or_, disjunctions)

[wence-]

Can I propose a rewrite like this, which I think separates the handler logic from the conjunction/disjunction a little more clearly? WDYT?

(Probably needs some minor modifications for py 3.9 compat (with the type-annotations)).

from functools import reduce, partial
import operator
import numpy as np


def handle_in(column, value):
    if not isinstance(value, (list, set, tuple)):
        raise TypeError("Value of 'in' filter must be a " "list, set, or tuple.")
    return reduce(operator.or_, (operator.eq(column, value) for v in value))


def handle_is(column, value, *, negate):
    if value not in {np.nan, None}:
        raise TypeError("Value of 'is' or 'is not' filter " "must be np.nan or None.")
    return ~column.isna() if negate else column.isna()


def _apply_post_filters(df, filters: list[tuple | list[tuple]]):
    # Apply DNF filters to an in-memory DataFrame
    #
    # Disjunctive normal form (DNF) means that the inner-most
    # tuple describes a single column predicate. These inner
    # predicates are combined with an AND conjunction into a
    # larger predicate. The outer-most list then combines all
    # of the combined filters with an OR disjunction.

    if not filters:
        # No filters to apply
        return df

    handlers = {
        "==": operator.eq,
        "!=": operator.ne,
        "<": operator.lt,
        "<=": operator.le,
        ">": operator.gt,
        ">=": operator.ge,
        "in": handle_in,
        "is": partial(handle_is, negate=False),
        "is not": partial(handle_is, negate=True),
    }

    try:
        # Disjunction loop
        #
        # All elements of `disjunctions` shall be combined with
        # an `OR` disjunction (operator.or_)
        disjunctions = []
        expressions = [f if isinstance(f, list) else [f] for f in filters]
        for expr in expressions:
            conjunction = reduce(
                operator.and_,
                (handlers[op](df[column], value) for (column, op, value) in expr),
            )
            disjunctions.append(conjunction)

        return reduce(operator.or_, disjunctions)
    except (KeyError, TypeError):
        warnings.warn(...)
        return df

[wence-]

The reduction over the disjunctions could be merged in (so no looping) but I think it's a little less readable (would be something like):

        return reduce(
            operator.or_,
            (
                reduce(
                    operator.and_,
                    (handlers[op](df[column], value) for (column, op, value) in expr),
                )
                for expr in expressions
            ),
        )

[rjzamora]

Thanks for the suggestion. I revised the code to look more like your example (there were some minor bugs, and I could't quite get the type annotations right yet - so left them out). Wouldn't mind going with the loop-free code, but didn't get a chance to try it yet.

[wence-]

Can you add some tests with more than two predicates being anded together (and being ored together)?

[GregoryKimball]

(Linking the discussion in #12512)

[rjzamora]

Small update: Removed the post_filters argument in favor of calling _apply_post_filters manually in dask_cudf.

[wence-]

I am not fully sure this is correct handling (my suggestion might also have been wrong).

AIUI, these are possible inputs as filters:

• (A, op, B) => (A op B)
• [(A, op, B)] => (A op B)
• [(A, op, B), (C, op, D)] => (A op B) v (C op D)
• [[(A, op, B), (C, op, D)], (E, op, F)] => ((A op B) ^ (C op D)) v (E op F)
• [[(A, op, B), (C, op, D)], [(E, op, F), (G, op H)]] => ((A op B) ^ (C op D)) v ((E op F) ^ (G op H))

So the input type is tuple | list[tuple | list[tuple]]

But this code only handles tuple | list[list[tuple]].

TBF, my code only handled list[tuple | list[tuple]].

To rephrase, who should do the sanitisation of the filters argument to this function? It would be much easier if, by the time we got here, we always just had list[list[tuple]]. That sanitisation could either be pushed up to read_parquet or else here but a little bit earlier, so we would say something like:

def _apply_filters(df, filters : tuple | list[tuple | list[tuple]]):
    if isinstance(filters, tuple): # singleton (A op B)
        filters = [filters]
    filters = [f if isinstance(f, list) else [f] for f in filters]
    ...
    for expr in filters: # Now everything is pre-processed correctly.

[wence-]

Mostly commentary, but a few minor cleanups.

[wence-]

Oops, sorry!

[wence-]

Is it easier to accept empty containers as "no filters" and normalise them to None (rather than requiring specifically None as empty filters).

So:

if filters:
   ... validate
else:
   return None

[wence-]

OK, so now we definitively have a list-of-lists.

[wence-]

OK, so each entry must be a non-empty list.

[wence-]

And each entry in that non-empty list must be a 3-tuple of appropriate type.

[wence-]

tl;dr: nothing to do here.

Could write:

if not all(isinstance(item, tuple) and len(item) == 3 for item in conjunction):
   raise TypeError("Every predicate must be ...")

But I guess then it's hard to point at the bad one, unless one abuses the walrus operator like so:

if not all(isinstance(item := pred, tuple) and len(pred) == 3 for pred in conjunctions):
   raise TypeError(... item)

Which is kind of ugly.

[rjzamora]

Right - I don't expect the number of predicates to get very large here. It seems reasonable to just call _validate_predicate in a loop.

[rjzamora]

Thanks for the the reviews @wence- @bdice ! This should be good for a final look when you have time.

[bdice]

Thanks @rjzamora. I learned a thing or two from this review! Appreciate your effort here.

[wence-]

Thanks @rjzamora !

[rjzamora]

/merge

[shwina]

@rjzamora do you expect the work done here to be roughly the same for orc?

[rjzamora]

do you expect the work done here to be roughly the same for orc?

Yes, the _apply_post_filters code added here should be independent of the IO mechanism, since it is a post-IO filtering operation. We could easily move it out of the parquet.py file and use it for ORC (and even CSV).

[shwina]

Thanks, I'll go ahead and open a PR to do the same for ORC then!