From 8875c76dd8b5b5f8d7b737d4ee5175c85b672808 Mon Sep 17 00:00:00 2001 From: Mohak Gupta Date: Fri, 4 Sep 2026 09:34:03 +0530 Subject: [PATCH 1/3] Add Java/JNI applyNullMask to ColumnView Applies a BOOL8 mask column onto another column, nulling out any row where the mask is false or null. Built on bools_to_mask plus superimpose_and_sanitize_nulls so it handles nested (list/struct) columns correctly, per ttnghia's comment on the issue. Closes #16764 --- .../main/java/ai/rapids/cudf/ColumnView.java | 20 +++++++++++++ java/src/main/native/src/ColumnViewJni.cpp | 27 +++++++++++++++++ .../test/java/ai/rapids/cudf/IfElseTest.java | 29 ++++++++++++++++++- 3 files changed, 75 insertions(+), 1 deletion(-) diff --git a/java/src/main/java/ai/rapids/cudf/ColumnView.java b/java/src/main/java/ai/rapids/cudf/ColumnView.java index 269a28ad8549..77b312507995 100644 --- a/java/src/main/java/ai/rapids/cudf/ColumnView.java +++ b/java/src/main/java/ai/rapids/cudf/ColumnView.java @@ -630,6 +630,24 @@ public final ColumnVector ifElse(Scalar trueValue, Scalar falseValue) { return new ColumnVector(result); } + /** + * Returns a copy of this column with each row set to null wherever the corresponding row + * in booleanMask is false or null, leaving the other rows as they already are (including + * any that were already null). + *

+ * This is a convenience over {@code ifElse}, which otherwise needs a null scalar of this + * column's type constructed just to null out the false rows. + * @param booleanMask a BOOL8 column with the same row count as this column + * @return a new column with this column's values, nulled out where booleanMask is not true + */ + public final ColumnVector applyNullMask(ColumnView booleanMask) { + if (!booleanMask.getType().equals(DType.BOOL8)) { + throw new IllegalArgumentException("Mask column must be of type BOOL8, found " + + booleanMask.getType()); + } + return new ColumnVector(applyNullMask(getNativeView(), booleanMask.getNativeView())); + } + ///////////////////////////////////////////////////////////////////////////// // Slice/Split and Concatenate ///////////////////////////////////////////////////////////////////////////// @@ -5128,6 +5146,8 @@ private static native long scan(long viewHandle, long aggregation, private static native long ifElseSS(long predVec, long trueScalar, long falseScalar) throws CudfException; + private static native long applyNullMask(long baseHandle, long boolMaskHandle) throws CudfException; + private static native long reduce(long viewHandle, long aggregation, int dtype, int scale) throws CudfException; private static native long segmentedReduce(long dataViewHandle, long offsetsViewHandle, diff --git a/java/src/main/native/src/ColumnViewJni.cpp b/java/src/main/native/src/ColumnViewJni.cpp index fe1a46ae6e72..78d1bd0d309a 100644 --- a/java/src/main/native/src/ColumnViewJni.cpp +++ b/java/src/main/native/src/ColumnViewJni.cpp @@ -303,6 +303,33 @@ JNIEXPORT jlong JNICALL Java_ai_rapids_cudf_ColumnView_ifElseSS( JNI_CATCH(env, 0); } +JNIEXPORT jlong JNICALL Java_ai_rapids_cudf_ColumnView_applyNullMask(JNIEnv* env, + jclass, + jlong j_col, + jlong j_bool_mask) +{ + JNI_NULL_CHECK(env, j_col, "column is null", 0); + JNI_NULL_CHECK(env, j_bool_mask, "mask column is null", 0); + JNI_TRY + { + cudf::jni::auto_set_device(env); + auto const col_view = reinterpret_cast(j_col); + auto const mask_view = reinterpret_cast(j_bool_mask); + + auto [bool_mask, bool_null_count] = cudf::bools_to_mask(*mask_view); + auto copy = std::make_unique(*col_view); + auto result = cudf::structs::detail::superimpose_and_sanitize_nulls( + static_cast(bool_mask->data()), + bool_null_count, + std::move(copy), + cudf::get_default_stream(), + cudf::get_current_device_resource_ref()); + + return release_as_jlong(result); + } + JNI_CATCH(env, 0); +} + JNIEXPORT jlong JNICALL Java_ai_rapids_cudf_ColumnView_getElement(JNIEnv* env, jclass, jlong from, diff --git a/java/src/test/java/ai/rapids/cudf/IfElseTest.java b/java/src/test/java/ai/rapids/cudf/IfElseTest.java index 2b78d8ab369b..97467cde0cd5 100644 --- a/java/src/test/java/ai/rapids/cudf/IfElseTest.java +++ b/java/src/test/java/ai/rapids/cudf/IfElseTest.java @@ -1,6 +1,6 @@ /* * - * SPDX-FileCopyrightText: Copyright (c) 2020, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2020-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 * */ @@ -1167,4 +1167,31 @@ void testMismatchedTypesSS() { assertThrows(CudfException.class, () -> pred.ifElse(trueScalar, falseScalar)); } } + + @Test + void testApplyNullMask() { + try (ColumnVector input = ColumnVector.fromBoxedInts(0, 100, 1, 2, Integer.MIN_VALUE, null); + ColumnVector mask = ColumnVector.fromBoxedBooleans(true, false, true, null, false, true); + ColumnVector result = input.applyNullMask(mask); + ColumnVector expected = ColumnVector.fromBoxedInts(0, null, 1, null, null, null)) { + assertColumnsAreEqual(expected, result); + } + } + + @Test + void testApplyNullMaskAllTrueIsValuePreserving() { + try (ColumnVector input = ColumnVector.fromBoxedInts(0, 100, null, 2); + ColumnVector mask = ColumnVector.fromBoxedBooleans(true, true, true, true); + ColumnVector result = input.applyNullMask(mask)) { + assertColumnsAreEqual(input, result); + } + } + + @Test + void testApplyNullMaskRejectsNonBooleanMask() { + try (ColumnVector input = ColumnVector.fromBoxedInts(0, 100, 1, 2); + ColumnVector mask = ColumnVector.fromBoxedInts(1, 0, 1, 0)) { + assertThrows(IllegalArgumentException.class, () -> input.applyNullMask(mask)); + } + } } From 1114b230475fbf624a8c49b71c6b222cc48edb5e Mon Sep 17 00:00:00 2001 From: Mohak Gupta Date: Sat, 5 Sep 2026 11:20:25 +0530 Subject: [PATCH 2/3] Validate applyNullMask row counts match before calling native CodeRabbit flagged that a mismatched mask size would reach the native side as a bitmask sized for the wrong row count instead of a clean error. Check row counts in Java first, same place the BOOL8 type is already checked. Signed-off-by: Mohak Gupta --- java/src/main/java/ai/rapids/cudf/ColumnView.java | 4 ++++ java/src/test/java/ai/rapids/cudf/IfElseTest.java | 8 ++++++++ 2 files changed, 12 insertions(+) diff --git a/java/src/main/java/ai/rapids/cudf/ColumnView.java b/java/src/main/java/ai/rapids/cudf/ColumnView.java index 77b312507995..a33123a37bb0 100644 --- a/java/src/main/java/ai/rapids/cudf/ColumnView.java +++ b/java/src/main/java/ai/rapids/cudf/ColumnView.java @@ -645,6 +645,10 @@ public final ColumnVector applyNullMask(ColumnView booleanMask) { throw new IllegalArgumentException("Mask column must be of type BOOL8, found " + booleanMask.getType()); } + if (booleanMask.getRowCount() != getRowCount()) { + throw new IllegalArgumentException("Mask column row count (" + booleanMask.getRowCount() + + ") does not match this column's row count (" + getRowCount() + ")"); + } return new ColumnVector(applyNullMask(getNativeView(), booleanMask.getNativeView())); } diff --git a/java/src/test/java/ai/rapids/cudf/IfElseTest.java b/java/src/test/java/ai/rapids/cudf/IfElseTest.java index 97467cde0cd5..4fc46ac2dcad 100644 --- a/java/src/test/java/ai/rapids/cudf/IfElseTest.java +++ b/java/src/test/java/ai/rapids/cudf/IfElseTest.java @@ -1194,4 +1194,12 @@ void testApplyNullMaskRejectsNonBooleanMask() { assertThrows(IllegalArgumentException.class, () -> input.applyNullMask(mask)); } } + + @Test + void testApplyNullMaskRejectsRowCountMismatch() { + try (ColumnVector input = ColumnVector.fromBoxedInts(0, 100, 1, 2); + ColumnVector mask = ColumnVector.fromBoxedBooleans(true, false, true)) { + assertThrows(IllegalArgumentException.class, () -> input.applyNullMask(mask)); + } + } } From 341f7386c6ae6d9b628cb67a99095eb351ddfe18 Mon Sep 17 00:00:00 2001 From: Mohak Gupta Date: Sat, 5 Sep 2026 21:39:12 +0530 Subject: [PATCH 3/3] Add nested column coverage for applyNullMask CodeRabbit flagged that only primitive-column cases were tested, even though superimpose_and_sanitize_nulls was chosen specifically to handle nested (list/struct) columns correctly per ttnghia's comment on the issue. Adds a struct-children propagation test and a list-offset-purge test mirroring the existing mergeAndSetValidity coverage for the same scenarios. Signed-off-by: Mohak Gupta --- .../test/java/ai/rapids/cudf/IfElseTest.java | 63 +++++++++++++++++++ 1 file changed, 63 insertions(+) diff --git a/java/src/test/java/ai/rapids/cudf/IfElseTest.java b/java/src/test/java/ai/rapids/cudf/IfElseTest.java index 4fc46ac2dcad..63db67696b47 100644 --- a/java/src/test/java/ai/rapids/cudf/IfElseTest.java +++ b/java/src/test/java/ai/rapids/cudf/IfElseTest.java @@ -12,10 +12,14 @@ import org.junit.jupiter.params.provider.Arguments; import org.junit.jupiter.params.provider.MethodSource; +import java.util.Arrays; import java.util.stream.Stream; import static ai.rapids.cudf.AssertUtils.assertColumnsAreEqual; +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertFalse; import static org.junit.jupiter.api.Assertions.assertThrows; +import static org.junit.jupiter.api.Assertions.assertTrue; public class IfElseTest extends CudfTestBase { private static Stream createBooleanVVParams() { @@ -1202,4 +1206,63 @@ void testApplyNullMaskRejectsRowCountMismatch() { assertThrows(IllegalArgumentException.class, () -> input.applyNullMask(mask)); } } + + @Test + void testApplyNullMaskPropagatesToStructChildren() { + try (ColumnVector c0 = ColumnVector.fromInts(1, 2, 3, 4, 5); + ColumnVector c1 = ColumnVector.fromInts(10, 20, 30, 40, 50); + ColumnVector struct = ColumnVector.makeStruct(c0, c1); + ColumnVector mask = ColumnVector.fromBoxedBooleans(true, true, false, null, true); + ColumnVector result = struct.applyNullMask(mask); + HostColumnVector hostResult = result.copyToHost()) { + assertEquals(2, hostResult.getNullCount(), "parent null count"); + assertFalse(hostResult.isNull(0)); + assertFalse(hostResult.isNull(1)); + assertTrue(hostResult.isNull(2)); + assertTrue(hostResult.isNull(3)); + assertFalse(hostResult.isNull(4)); + + // Each child should have the same null mask as the parent. + assertEquals(2, hostResult.getNumChildren()); + for (int i = 0; i < hostResult.getNumChildren(); i++) { + HostColumnVectorCore child = hostResult.getChildColumnView(i); + assertEquals(2, child.getNullCount(), "child " + i + " null count"); + assertTrue(child.isNull(2), "child " + i + " row 2"); + assertTrue(child.isNull(3), "child " + i + " row 3"); + } + } + } + + @Test + void testApplyNullMaskPurgesListOffsetsOfMaskedRows() { + HostColumnVector.DataType intType = new HostColumnVector.BasicType(true, DType.INT32); + HostColumnVector.DataType listType = new HostColumnVector.ListType(true, intType); + try (ColumnVector list = ColumnVector.fromLists(listType, + Arrays.asList(1, 2), + Arrays.asList(3, 4, 5), + Arrays.asList(6), // will be masked null. + Arrays.asList(7, 8, 9, 10), // will be masked null. + Arrays.asList(11)); + ColumnVector mask = ColumnVector.fromBoxedBooleans(true, true, false, null, true); + ColumnVector result = list.applyNullMask(mask); + HostColumnVector hostResult = result.copyToHost()) { + assertEquals(2, hostResult.getNullCount(), "parent null count"); + assertTrue(hostResult.isNull(2)); + assertTrue(hostResult.isNull(3)); + + // Rows 2 and 3 collapse so the inner INT should have only 6 elements. + assertEquals(1, hostResult.getNumChildren()); + HostColumnVectorCore intChild = hostResult.getChildColumnView(0); + assertEquals(6, intChild.getRowCount(), "purged inner row count"); + int[] expectedInner = {1, 2, 3, 4, 5, 11}; + for (int i = 0; i < expectedInner.length; i++) { + assertEquals(expectedInner[i], intChild.getInt(i), "inner " + i); + } + HostMemoryBuffer offsets = hostResult.getOffsets(); + int[] expectedOffsets = {0, 2, 5, 5, 5, 6}; + for (int i = 0; i < expectedOffsets.length; i++) { + assertEquals(expectedOffsets[i], offsets.getInt(i * 4L), "offset " + i); + } + } + } }