-
Notifications
You must be signed in to change notification settings - Fork 4.2k
GH-49376: [Python][Parquet] Add ability to write Bloom filters from pyarrow #49377
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
Changes from 5 commits
a9ca4d7
c0fb063
176fe77
c74fa69
c67af90
1b39f04
965729a
File filter
Filter by extension
Conversations
Jump to
Diff view
Diff view
There are no files selected for viewing
| Original file line number | Diff line number | Diff line change |
|---|---|---|
|
|
@@ -51,6 +51,10 @@ cimport cpython as cp | |
| _DEFAULT_ROW_GROUP_SIZE = 1024*1024 | ||
| _MAX_ROW_GROUP_SIZE = 64*1024*1024 | ||
|
|
||
| # from definition of BloomFilterOptions struct | ||
| _DEFAULT_BLOOM_FILTER_NDV = 1024*1024 | ||
| _DEFAULT_BLOOM_FILTER_FPP = 0.05 | ||
|
|
||
|
|
||
| cdef Type _unwrap_list_type(obj) except *: | ||
| if obj is ListType: | ||
|
|
@@ -1992,13 +1996,15 @@ cdef shared_ptr[WriterProperties] _create_writer_properties( | |
| write_page_checksum=False, | ||
| sorting_columns=None, | ||
| store_decimal_as_integer=False, | ||
| use_content_defined_chunking=False) except *: | ||
| use_content_defined_chunking=False, | ||
| bloom_filter_options=None) except *: | ||
|
|
||
| """General writer properties""" | ||
| cdef: | ||
| shared_ptr[WriterProperties] properties | ||
| WriterProperties.Builder props | ||
| CdcOptions cdc_options | ||
| BloomFilterOptions bloom_opts | ||
|
|
||
| # data_page_version | ||
|
|
||
|
|
@@ -2122,6 +2128,48 @@ cdef shared_ptr[WriterProperties] _create_writer_properties( | |
| raise TypeError( | ||
| "'column_encoding' should be a dictionary or a string") | ||
|
|
||
| # bloom filters | ||
| if bloom_filter_options is not None: | ||
| if isinstance(bloom_filter_options, dict): | ||
|
Member
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. Can we factor this out into a helper function? This function is becoming much too large IMHO. |
||
| # for each entry in bloom_filter_options, {"path": {"ndv": ndv, "fpp", fpp}} | ||
| # convert (ndv,fpp) to BloomFilterOptions struct and pass to props | ||
| for column, _bloom_opts in bloom_filter_options.items(): | ||
| # set defaults | ||
| bloom_opts.ndv = _DEFAULT_BLOOM_FILTER_NDV | ||
| bloom_opts.fpp = _DEFAULT_BLOOM_FILTER_FPP | ||
| if isinstance(_bloom_opts, dict): | ||
| if "ndv" in _bloom_opts: | ||
| ndv = _bloom_opts["ndv"] | ||
| if isinstance(ndv, int): | ||
| if ndv < 0: | ||
|
Member
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. ndv <= 0?
Contributor
Author
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. I'm not sure if NDV can be 0 or not...I can't seem to find validation in either the C++ nor rust implementations. This test just ensures it's a positive value. Honestly, I don't see why you'd need a bloom filter for a column with no values, so I'm fine following your suggestion. |
||
| raise ValueError( | ||
| f"'ndv' for column '{column}' must be positive, got {ndv}") | ||
| bloom_opts.ndv = ndv | ||
| else: | ||
| raise TypeError( | ||
| f"'ndv' for column '{column}' must be an int") | ||
|
Member
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. should we notice they're in bloom filter options? |
||
| if "fpp" in _bloom_opts: | ||
| fpp = _bloom_opts["fpp"] | ||
| if isinstance(fpp, float): | ||
|
Member
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. I don't know whether casting to
Member
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. I don't think so. We don't want to accept strings, for example. |
||
| if fpp <= 0.0 or fpp >= 1.0: | ||
| raise ValueError( | ||
| f"'fpp' for column '{column}' must be in (0.0, 1,0), got {fpp}") | ||
|
Member
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. 1,0 typo here? |
||
| bloom_opts.fpp = fpp | ||
| else: | ||
| raise TypeError( | ||
| f"'fpp' for column '{column}' must be a float") | ||
| elif isinstance(_bloom_opts, bool): | ||
| if not _bloom_opts: | ||
|
mapleFU marked this conversation as resolved.
Outdated
|
||
| props.disable_bloom_filter(tobytes(column)) | ||
| continue | ||
|
Member
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. I understand that bloom filter is disabled by default but should we be explicit here calling
Contributor
Author
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. Yes, I agree it would be better to be explicit here. Thanks! |
||
| else: | ||
| raise TypeError( | ||
| f"'bloom_filter_options:{column}' must be a boolean or a dictionary") | ||
|
|
||
| props.enable_bloom_filter(tobytes(column), bloom_opts) | ||
| else: | ||
| raise TypeError("'bloom_filter_options' must be a dictionary") | ||
|
|
||
| # size limits | ||
| if data_page_size is not None: | ||
| props.data_pagesize(data_page_size) | ||
|
|
@@ -2317,7 +2365,8 @@ cdef class ParquetWriter(_Weakrefable): | |
| sorting_columns=None, | ||
| store_decimal_as_integer=False, | ||
| use_content_defined_chunking=False, | ||
| write_time_adjusted_to_utc=False): | ||
| write_time_adjusted_to_utc=False, | ||
| bloom_filter_options=None): | ||
| cdef: | ||
| shared_ptr[WriterProperties] properties | ||
| shared_ptr[ArrowWriterProperties] arrow_properties | ||
|
|
@@ -2353,7 +2402,8 @@ cdef class ParquetWriter(_Weakrefable): | |
| write_page_checksum=write_page_checksum, | ||
| sorting_columns=sorting_columns, | ||
| store_decimal_as_integer=store_decimal_as_integer, | ||
| use_content_defined_chunking=use_content_defined_chunking | ||
| use_content_defined_chunking=use_content_defined_chunking, | ||
| bloom_filter_options=bloom_filter_options | ||
| ) | ||
| arrow_properties = _create_arrow_writer_properties( | ||
| use_deprecated_int96_timestamps=use_deprecated_int96_timestamps, | ||
|
|
||
Uh oh!
There was an error while loading. Please reload this page.