diff --git a/cpp/src/io/parquet/bloom_filter_reader.cu b/cpp/src/io/parquet/bloom_filter_reader.cu index f222365de180..d2d7fcac9597 100644 --- a/cpp/src/io/parquet/bloom_filter_reader.cu +++ b/cpp/src/io/parquet/bloom_filter_reader.cu @@ -402,10 +402,12 @@ size_t aggregate_reader_metadata::get_bloom_filter_alignment() const // Required alignment: // https://github.com/NVIDIA/cuCollections/blob/deab5799f3e4226cb8a49acf2199c03b14941ee4/include/cuco/detail/bloom_filter/bloom_filter_impl.cuh#L55-L67 using policy_type = cuco::arrow_filter_policy; - return alignof(cuco::bloom_filter_ref, - cuco::thread_scope_thread, - policy_type>::filter_block_type); + auto constexpr alignment = alignof(cuco::bloom_filter_ref, + cuco::thread_scope_thread, + policy_type>::filter_block_type); + static_assert((alignment & (alignment - 1)) == 0, "Alignment must be a power of 2"); + return std::max(alignment, rmm::CUDA_ALLOCATION_ALIGNMENT); } std::vector aggregate_reader_metadata::read_bloom_filters( diff --git a/cpp/tests/io/experimental/hybrid_scan_test.cpp b/cpp/tests/io/experimental/hybrid_scan_test.cpp index 6ec572d1e493..5cb787f451ee 100644 --- a/cpp/tests/io/experimental/hybrid_scan_test.cpp +++ b/cpp/tests/io/experimental/hybrid_scan_test.cpp @@ -31,9 +31,10 @@ #include #include +#include #include -auto constexpr bloom_filter_alignment = 32; +auto constexpr bloom_filter_alignment = rmm::CUDA_ALLOCATION_ALIGNMENT; namespace { diff --git a/python/cudf/cudf/tests/data/parquet/bloom_filter_alignment.parquet b/python/cudf/cudf/tests/data/parquet/bloom_filter_alignment.parquet new file mode 100644 index 000000000000..26441e59257a Binary files /dev/null and b/python/cudf/cudf/tests/data/parquet/bloom_filter_alignment.parquet differ diff --git a/python/cudf/cudf/tests/test_parquet.py b/python/cudf/cudf/tests/test_parquet.py index d1e82a552ad0..94eb2c794a58 100644 --- a/python/cudf/cudf/tests/test_parquet.py +++ b/python/cudf/cudf/tests/test_parquet.py @@ -4523,6 +4523,47 @@ def test_parquet_bloom_filters( ) +@pytest.fixture(params=["cuda", "pool", "cuda_async"]) +def memory_resource(request): + import rmm + + current_mr = rmm.mr.get_current_device_resource() + + kind = request.param + if kind == "cuda": + mr = rmm.mr.CudaMemoryResource() + elif kind == "pool": + base = rmm.mr.CudaMemoryResource() + free, _ = rmm.mr.available_device_memory() + size = int(round(free * 0.5 / 256) * 256) + mr = rmm.mr.PoolMemoryResource(base, size, size) + elif kind == "cuda_async": + mr = rmm.mr.CudaAsyncMemoryResource() + + rmm.mr.set_current_device_resource(mr) + + try: + yield mr + finally: + rmm.mr.set_current_device_resource(current_mr) + + +@pytest.mark.parametrize("columns", [["r_reason_desc"], None]) +def test_parquet_bloom_filters_alignment(datadir, columns, memory_resource): + fname = datadir / "bloom_filter_alignment.parquet" + filters = [("r_reason_desc", "==", "Did not like the color")] + + # Read expected table using pyarrow + expected = pq.read_table(fname, columns=columns, filters=filters) + + # Read with cudf using the memory resource from fixture + read = cudf.read_parquet( + fname, columns=columns, filters=filters + ).to_arrow() + + assert_eq(expected, read) + + def test_parquet_reader_unsupported_compression(datadir): fname = datadir / "hadoop_lz4_compressed.parquet"