Merge branch 'main' into docs-libcudf-introspect-overflow

Author: PointKernel

cpp/libcudf_streaming/src/detail/device_bloom_filter.cu

@@ -58,7 +58,16 @@ using BloomFilterRefType =
   cuco::bloom_filter_ref<KeyType,
                          cuco::extent<std::size_t>,
                          cuco::thread_scope_device,
-                         cuco::arrow_filter_policy<KeyType, cuco::identity_hash>>;
+                         cuco::parametric_filter_policy<cuco::identity_hash<KeyType>,
+                                                        std::uint32_t,
+                                                        8,
+                                                        8,
+                                                        8,
+                                                        1,
+                                                        1,
+                                                        8,
+                                                        false,
+                                                        false>>;
 using StorageType = BloomFilterRefType::filter_block_type;
 
 }  // namespace

cpp/src/io/parquet/arrow_filter_policy.cuh

@@ -1,9 +1,8 @@
 /*
- * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION.
+ * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
  * SPDX-License-Identifier: Apache-2.0
  */
 
-#include "arrow_filter_policy.cuh"
 #include "compact_protocol_reader.hpp"
 #include "expression_transform_helpers.hpp"
 #include "io/utilities/time_utils.hpp"
@@ -25,6 +24,7 @@
 #include <rmm/device_buffer.hpp>
 #include <rmm/exec_policy.hpp>
 
+#include <cuco/bloom_filter_policies.cuh>
 #include <cuco/bloom_filter_ref.cuh>
 #include <cuda/iterator>
 #include <thrust/tabulate.h>
@@ -36,6 +36,30 @@
 namespace cudf::io::parquet::detail {
 namespace {
 
+/**
+ * @brief Policy describing the Apache Arrow Block-Split Bloom Filter, hashing keys with cudf's
+ * `XXHash_64` (so that `cudf::string_view` and other cudf types are hashed by content, matching the
+ * Apache Parquet/Arrow bloom filter specification).
+ *
+ * Uses cuco's `parametric_filter_policy` with the Apache Arrow layout: 256-bit blocks (8 x
+ * `uint32_t`), 8 fingerprint bits per key, fully horizontal add (Theta=8) and fully vertical
+ * contains (Phi=8). This layout is bit-compatible with Apache Arrow, as verified by cuCollections
+ * `tests/bloom_filter/arrow_compat_test.cu`.
+ *
+ * @tparam Key The type of the values to generate a fingerprint for.
+ */
+template <class Key>
+using arrow_filter_policy = cuco::parametric_filter_policy<cudf::hashing::detail::XXHash_64<Key>,
+                                                           std::uint32_t,
+                                                           8,
+                                                           8,
+                                                           8,
+                                                           1,
+                                                           1,
+                                                           8,
+                                                           false,
+                                                           false>;
+
 /**
  * @brief Converts bloom filter membership results (for each column chunk) to a device column.
  *

cpp/src/io/parquet/bloom_filter_reader.cu

@@ -1,5 +1,5 @@
 /*
- * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION.
+ * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
  * SPDX-License-Identifier: Apache-2.0
  */
 #pragma once
@@ -165,8 +165,7 @@ using storage_ref_type =
   cuco::bucket_storage_ref<mark_key_type, mark_join_bucket_size, cuco::extent<std::size_t>>;
 using right_key_type = cuco::pair<hash_value_type, rhs_index_type>;
 
-using bloom_filter_policy_type =
-  cuco::default_filter_policy<cuco::detail::identity_hash<hash_value_type>, hash_value_type, 2U>;
+using bloom_filter_policy_type    = cuco::default_filter_policy<hash_value_type>;
 using bloom_filter_allocator_type = rmm::mr::polymorphic_allocator<cuda::std::byte>;
 using bloom_filter_type           = cuco::bloom_filter<hash_value_type,
                                                        cuco::extent<std::size_t>,

cpp/src/join/mark_join.cuh

@@ -1,10 +1,8 @@
 /*
- * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION.
+ * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
  * SPDX-License-Identifier: Apache-2.0
  */
 
-#include "src/io/parquet/arrow_filter_policy.cuh"
-
 #include <cudf_test/base_fixture.hpp>
 #include <cudf_test/column_utilities.hpp>
 #include <cudf_test/column_wrapper.hpp>
@@ -18,15 +16,29 @@
 #include <rmm/mr/polymorphic_allocator.hpp>
 
 #include <cuco/bloom_filter.cuh>
+#include <cuco/bloom_filter_policies.cuh>
+
+#include <cstdint>
 
 using StringType = cudf::string_view;
 
 class ParquetBloomFilterTest : public cudf::test::BaseFixture {};
 
 TEST_F(ParquetBloomFilterTest, TestStrings)
 {
-  using key_type    = StringType;
-  using policy_type = cudf::io::parquet::detail::arrow_filter_policy<key_type>;
+  using key_type = StringType;
+  // Apache Arrow Block-Split Bloom Filter layout, hashing keys with cudf's `XXHash_64` (matching
+  // `cudf::io::parquet::detail::arrow_filter_policy`).
+  using policy_type = cuco::parametric_filter_policy<cudf::hashing::detail::XXHash_64<key_type>,
+                                                     std::uint32_t,
+                                                     8,
+                                                     8,
+                                                     8,
+                                                     1,
+                                                     1,
+                                                     8,
+                                                     false,
+                                                     false>;
   using word_type   = policy_type::word_type;
 
   std::size_t constexpr num_filter_blocks = 4;

cpp/tests/io/parquet_bloom_filter_test.cu

@@ -2,17 +2,23 @@
 # SPDX-License-Identifier: Apache-2.0
 from __future__ import annotations
 
+import operator
+from functools import partial
 from typing import TYPE_CHECKING
 
 from numba_cuda_mlir import types
 from numba_cuda_mlir._mlir import ir as mlir_ir
-from numba_cuda_mlir._mlir.dialects import llvm
-from numba_cuda_mlir.extending import lowering_registry
+from numba_cuda_mlir._mlir.dialects import arith, llvm
+from numba_cuda_mlir.extending import lower_cast, lowering_registry
 from numba_cuda_mlir.lowering_utilities import convert
 from numba_cuda_mlir.models import PrimitiveModel, register_model
 
 from cudf.core.udf.api import Masked
-from cudf.core.udf.mlir_backend.masked_typing import MaskedType
+from cudf.core.udf.mlir_backend.masked_typing import (
+    MaskedType,
+    NAType,
+    na_type,
+)
 
 if TYPE_CHECKING:
     from numba_cuda_mlir.mlir_lowering import MLIRLower
@@ -62,6 +68,21 @@ def __init__(
         super().__init__(data_model_manager, masked_type, struct_type)
 
 
+def _extract_masked_value_valid(struct_val, value_mlir_ty, valid_ty):
+    """Pull the ``(value, valid)`` SSA values out of a ``Masked`` struct."""
+    v = llvm.extractvalue(
+        res=value_mlir_ty,
+        container=struct_val,
+        position=mlir_ir.DenseI64ArrayAttr.get([0]),
+    )
+    valid = llvm.extractvalue(
+        res=valid_ty,
+        container=struct_val,
+        position=mlir_ir.DenseI64ArrayAttr.get([1]),
+    )
+    return v, valid
+
+
 def _lower_masked_constructor(
     builder: MLIRLower, target: Var, args: list[Var], kwargs: list
 ) -> None:
@@ -105,6 +126,63 @@ def _lower_masked_getattr(
     builder.store_var(target, convert(field_value, target_mlir_ty))
 
 
+# ``cast(NA -> Masked)`` and ``cast(scalar -> Masked)``. Both build a Masked
+# struct for the target's value type; they differ only in the payload (NA has
+# none, so use undef) and the validity bit (NA -> invalid, scalar -> valid).
+# Triggered by branch unification, e.g. ``return x if cond else cudf.NA`` or
+# ``return 5``.
+def _cast_to_masked(context, builder, from_ty, to_ty, val):
+    value_mlir_ty = builder.get_mlir_type(to_ty.value_type)
+    if isinstance(from_ty, NAType):
+        value = llvm.UndefOp(value_mlir_ty)
+        valid = 0
+    else:
+        value = convert(val, value_mlir_ty)
+        valid = 1
+    valid_const = arith.constant(
+        result=builder.get_mlir_type(types.boolean), value=valid
+    )
+    return _pack_masked(builder, to_ty, value, valid_const)
+
+
+# ``cast(Masked -> Masked)``: branch unification across different
+# inner widths (e.g. one branch returns Masked(int32), another
+# returns Masked(float64); Numba unifies to Masked(float64)).
+# Promote the payload, preserve the validity bit.
+def _cast_masked_to_masked(context, builder, from_ty, to_ty, val):
+    if from_ty.value_type == to_ty.value_type:
+        return val
+    st = llvm.StructType(val.type)
+    m_val, m_valid = _extract_masked_value_valid(val, st.body[0], st.body[1])
+    value_mlir_ty = builder.get_mlir_type(to_ty.value_type)
+    casted = convert(m_val, value_mlir_ty)
+    return _pack_masked(builder, to_ty, casted, m_valid)
+
+
+# ``is``/``is not`` against NA are registered for both operand orders
+# (``m is NA`` and ``NA is m``), so find the MaskedType operand rather than
+# assuming which position it is in.
+def _masked_operand(builder, args):
+    """Return the ``MaskedType`` operand of a ``Masked``/``NA`` comparison."""
+    for arg in args:
+        if isinstance(builder.get_numba_type(arg.name), MaskedType):
+            return arg
+    raise TypeError("expected a MaskedType operand")
+
+
+# ``is``/``is not`` against NA both reduce to the validity bit: ``m is NA`` ->
+# ``not m.valid`` and ``m is not NA`` -> ``m.valid``. Registered for both
+# operators (and operand orders) via partials below.
+def _lower_masked_na_compare(builder, target, args, kwargs, *, is_null):
+    m = builder.load_var(_masked_operand(builder, args))
+    st = llvm.StructType(m.type)
+    _, valid = _extract_masked_value_valid(m, st.body[0], st.body[1])
+    if is_null:
+        one = arith.constant(valid.type, 1)
+        valid = arith.xori(valid, one)
+    builder.store_var(target, valid)
+
+
 def _register() -> None:
     """Register the data model and lowerings with ``numba_cuda_mlir``.
 
@@ -123,5 +201,17 @@ def _register() -> None:
 
     lowering_registry.lower_getattr_generic(MaskedType)(_lower_masked_getattr)
 
+    lower_cast(na_type, MaskedType)(_cast_to_masked)
+    for _scalar_cls in (types.Integer, types.Float, types.Boolean):
+        lower_cast(_scalar_cls, MaskedType)(_cast_to_masked)
+    lower_cast(MaskedType, MaskedType)(_cast_masked_to_masked)
+
+    is_na = partial(_lower_masked_na_compare, is_null=True)
+    is_not_na = partial(_lower_masked_na_compare, is_null=False)
+    lower(operator.is_, MaskedType, NAType)(is_na)
+    lower(operator.is_, NAType, MaskedType)(is_na)
+    lower(operator.is_not, MaskedType, NAType)(is_not_na)
+    lower(operator.is_not, NAType, MaskedType)(is_not_na)
+
 
 _register()