diff --git a/java/src/main/java/ai/rapids/cudf/ColumnView.java b/java/src/main/java/ai/rapids/cudf/ColumnView.java index 269a28ad854..a33123a37bb 100644 --- a/java/src/main/java/ai/rapids/cudf/ColumnView.java +++ b/java/src/main/java/ai/rapids/cudf/ColumnView.java @@ -630,6 +630,28 @@ 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()); + } + 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())); + } + ///////////////////////////////////////////////////////////////////////////// // Slice/Split and Concatenate ///////////////////////////////////////////////////////////////////////////// @@ -5128,6 +5150,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 fe1a46ae6e7..78d1bd0d309 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 2b78d8ab369..63db67696b4 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 * */ @@ -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() { @@ -1167,4 +1171,98 @@ 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)); + } + } + + @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)); + } + } + + @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); + } + } + } }