Repository: rapidsai/cudf
Pull request #18845: Null-handling for Transforms
Author: login: lamarrr
Labels: name: feature request
description: New feature or request
name: libcudf
description: Affects libcudf (C++/CUDA) code.
name: Java
description: Affects Java cuDF API.
name: breaking
description: Breaking change
Created: 2025-05-15T18:24:52Z | Merged: 2025-05-29T21:48:41Z

## Description

This merge request implements null-handling for Transforms. The implementation outputs a null row if any of its input rows are null. The output bitmask is computed in a single pass using the null-masks of the input columns before the transform function is run on the valid output rows.

It also handles the scalar input edge cases:

- Scalar is not null; implies it is valid across the base column's size. We also can't use its bitmask as it is a bitmask of size 1, otherwise, we'd encounter an out-of-bounds error. We exclude the bitmask from the `bitmask_and` operation.

- Scalar is null; implies it is null across the base column's size. This would make all the output rows null. The scalar's bitmask is also of size 1. In this case, we output a null-mask of the base column's size, and mark all the elements as null.


Follows up on https://github.com/rapidsai/cudf/issues/18023, https://github.com/rapidsai/cudf/issues/18820

<!-- Provide a standalone description of changes in this PR. -->
<!-- Reference any issues closed by this PR with "closes #1234". -->
<!-- Note: The pull request title will be included in the CHANGELOG. -->

## Checklist
- [x] I am familiar with the [Contributing Guidelines](https://github.com/rapidsai/cudf/blob/HEAD/CONTRIBUTING.md).
- [x] New or existing tests cover these changes.
- [x] The documentation is up to date with these changes.


Linked issue #1234: [FEA] Create a single RAPIDS library
**Is your feature request related to a problem? Please describe.**
Hi all,

I would like to share with you the following idea:

So just as we have the following imports:

> import pandas as pd
> import tensorflow as tf
> import matplotlib as plt

why do not create a rapids library that could be imported as follows?

> import rapids as rs

or

> from rapids import cudf

I am aware that there are pros about having each library separate, but I think there are also pros about having a single library, such us

- easier for the data scientists to install: conda install rapids
- it creates a kind of identity around rapids (we do not think on cudf or cuml as independent identities, but as a whole).
- we might adopt 'rs' as a shortcut for rapids (currently, I have only found the following use of 'rs' import rhinoscriptsyntax as  rs)

Please, share your thoughts about the above.

Thanks!

P.S.: There are also cons about having a huge library, but I think the pros exceed the cons.
Linked issue comments: user: login: kkraus14
created_at: 2019-07-03T22:56:24Z
body: This is out of scope for now, possibly to be revisited in the future.
author_association: COLLABORATOR

Diff:
diff --git a/cpp/include/cudf/jit/span.cuh b/cpp/src/jit/span.cuh
similarity index 65%
rename from cpp/include/cudf/jit/span.cuh
rename to cpp/src/jit/span.cuh
index 709bff3e627..1b07dc6feba 100644
--- a/cpp/include/cudf/jit/span.cuh
+++ b/cpp/src/jit/span.cuh
@@ -16,8 +16,9 @@
 #pragma once
 
 #include <cudf/types.hpp>
+#include <cudf/utilities/bit.hpp>
 
-namespace CUDF_EXPORT cudf {
+namespace cudf {
 
 namespace jit {
 
@@ -97,5 +98,59 @@ struct device_span {
   CUDF_HOST_DEVICE [[nodiscard]] constexpr element_type* end() const { return _data + _size; }
 };
 
+/**
+ * @brief A span type with optional/nullable elements.
+ *
+ * Optional implies the span contains nullable elements.
+ * The nullability of the elements is internally represented by an optional bitmask which can be
+ * nullptr when all the elements are non-null.
+ */
+template <typename T>
+struct device_optional_span : device_span<T> {
+ private:
+  using base               = device_span<T>;
+  bitmask_type* _null_mask = nullptr;
+
+ public:
+  CUDF_HOST_DEVICE constexpr device_optional_span() {}
+
+  /**
+   * @brief Constructs an optional span from a span and a null-mask.
+   *
+   * @param span Span containing the elements
+   * @param null_mask The null-mask determining the validity of the elements or nullptr if all
+   * valid.
+   */
+  CUDF_HOST_DEVICE device_optional_span(device_span<T> span, bitmask_type* null_mask)
+    : base{span}, _null_mask{null_mask}
+  {
+  }
+
+  /// @copydoc column_device_view::nullable
+  [[nodiscard]] CUDF_HOST_DEVICE bool nullable() const { return _null_mask != nullptr; }
+
+#ifdef __CUDACC__
+
+  /// @copydoc column_device_view::is_valid_nocheck
+  [[nodiscard]] __device__ bool is_valid_nocheck(size_t element_index) const
+  {
+    return bit_is_set(_null_mask, element_index);
+  }
+
+  /// @copydoc column_device_view::is_valid
+  [[nodiscard]] __device__ bool is_valid(size_t element_index) const
+  {
+    return not nullable() or is_valid_nocheck(element_index);
+  }
+
+  /// @copydoc column_device_view::is_null
+  [[nodiscard]] __device__ bool is_null(size_t element_index) const
+  {
+    return !is_valid(element_index);
+  }
+
+#endif
+};
+
 }  // namespace jit
-}  // namespace CUDF_EXPORT cudf
+}  // namespace cudf
diff --git a/cpp/src/transform/jit/kernel.cu b/cpp/src/transform/jit/kernel.cu
index e5797783c6a..7117b84cd9b 100644
--- a/cpp/src/transform/jit/kernel.cu
+++ b/cpp/src/transform/jit/kernel.cu
@@ -14,8 +14,9 @@
  * limitations under the License.
  */
 
+#include "jit/span.cuh"
+
 #include <cudf/column/column_device_view_base.cuh>
-#include <cudf/jit/span.cuh>
 #include <cudf/strings/string_view.cuh>
 #include <cudf/types.hpp>
 #include <cudf/utilities/traits.hpp>
@@ -60,6 +61,12 @@ struct column_accessor {
   {
     outputs[index].assign<T>(row, value);
   }
+
+  static __device__ bool is_null(cudf::mutable_column_device_view_core const* inputs,
+                                 cudf::size_type row)
+  {
+    return inputs[index].is_null(row);
+  }
 };
 
 template <typename T, int32_t Index>
@@ -67,17 +74,24 @@ struct span_accessor {
   using type                     = T;
   static constexpr int32_t index = Index;
 
-  static __device__ type& element(cudf::jit::device_span<T> const* spans, cudf::size_type row)
+  static __device__ type& element(cudf::jit::device_optional_span<T> const* spans,
+                                  cudf::size_type row)
   {
     return spans[index][row];
   }
 
-  static __device__ void assign(cudf::jit::device_span<T> const* outputs,
+  static __device__ void assign(cudf::jit::device_optional_span<T> const* outputs,
                                 cudf::size_type row,
                                 T value)
   {
     outputs[index][row] = value;
   }
+
+  static __device__ bool is_null(cudf::jit::device_optional_span<T> const* inputs,
+                                 cudf::size_type row)
+  {
+    return inputs[index].is_null(row);
+  }
 };
 
 template <typename Accessor>
@@ -86,19 +100,19 @@ struct scalar {
   static constexpr int32_t index = Accessor::index;
 
   static __device__ decltype(auto) element(cudf::mutable_column_device_view_core const* outputs,
-                                           cudf::size_type row)
+                                           cudf::size_type)
   {
     return Accessor::element(outputs, 0);
   }
 
   static __device__ decltype(auto) element(cudf::column_device_view_core const* inputs,
-                                           cudf::size_type row)
+                                           cudf::size_type)
   {
     return Accessor::element(inputs, 0);
   }
 
   static __device__ void assign(cudf::mutable_column_device_view_core const* outputs,
-                                cudf::size_type row,
+                                cudf::size_type,
                                 type value)
   {
     return Accessor::assign(outputs, 0, value);
@@ -121,6 +135,8 @@ CUDF_KERNEL void kernel(cudf::mutable_column_device_view_core const* outputs,
   thread_index_type const size   = outputs[0].size();
 
   for (auto i = start; i < size; i += stride) {
+    if (Out::is_null(outputs, i)) { continue; }
+
     if constexpr (has_user_data) {
       GENERIC_TRANSFORM_OP(user_data, i, &Out::element(outputs, i), In::element(inputs, i)...);
     } else {
@@ -145,6 +161,8 @@ CUDF_KERNEL void fixed_point_kernel(cudf::mutable_column_device_view_core const*
   for (auto i = start; i < size; i += stride) {
     typename Out::type result{numeric::scaled_integer<typename Out::type::rep>{0, output_scale}};
 
+    if (Out::is_null(outputs, i)) { continue; }
+
     if constexpr (has_user_data) {
       GENERIC_TRANSFORM_OP(user_data, i, &result, In::element(inputs, i)...);
     } else {
@@ -156,7 +174,7 @@ CUDF_KERNEL void fixed_point_kernel(cudf::mutable_column_device_view_core const*
 }
 
 template <bool has_user_data, typename Out, typename... In>
-CUDF_KERNEL void span_kernel(cudf::jit::device_span<typename Out::type> const* outputs,
+CUDF_KERNEL void span_kernel(cudf::jit::device_optional_span<typename Out::type> const* outputs,
                              cudf::column_device_view_core const* inputs,
                              void* user_data)
 {
@@ -168,6 +186,8 @@ CUDF_KERNEL void span_kernel(cudf::jit::device_span<typename Out::type> const* o
   thread_index_type const size   = outputs[0].size();
 
   for (auto i = start; i < size; i += stride) {
+    if (Out::is_null(outputs, i)) { continue; }
+
     if constexpr (has_user_data) {
       GENERIC_TRANSFORM_OP(user_data, i, &Out::element(outputs, i), In::element(inputs, i)...);
     } else {
diff --git a/cpp/src/transform/transform.cpp b/cpp/src/transform/transform.cpp
index 98c82429ad7..deb326ae050 100644
--- a/cpp/src/transform/transform.cpp
+++ b/cpp/src/transform/transform.cpp
@@ -16,6 +16,7 @@
 
 #include "jit/cache.hpp"
 #include "jit/parser.hpp"
+#include "jit/span.cuh"
 #include "jit/util.hpp"
 
 #include <cudf/column/column.hpp>
@@ -24,7 +25,6 @@
 #include <cudf/detail/nvtx/ranges.hpp>
 #include <cudf/detail/transform.hpp>
 #include <cudf/jit/runtime_support.hpp>
-#include <cudf/jit/span.cuh>
 #include <cudf/null_mask.hpp>
 #include <cudf/utilities/traits.hpp>
 #include <cudf/utilities/type_dispatcher.hpp>
@@ -38,6 +38,11 @@ namespace transformation {
 namespace jit {
 namespace {
 
+constexpr bool is_scalar(cudf::size_type base_column_size, cudf::size_type column_size)
+{
+  return column_size == 1 && column_size != base_column_size;
+}
+
 struct input_column_reflection {
   std::string type_name;
   bool is_scalar = false;
@@ -164,8 +169,8 @@ jitify2::Kernel get_kernel(std::string const& kernel_name, std::string const& cu
 
 input_column_reflection reflect_input_column(size_type base_column_size, column_view column)
 {
-  CUDF_EXPECTS(column.size() == 1 || column.size() == base_column_size, "");
-  return input_column_reflection{type_to_name(column.type()), column.size() != base_column_size};
+  return input_column_reflection{type_to_name(column.type()),
+                                 is_scalar(base_column_size, column.size())};
 }
 
 std::vector<input_column_reflection> reflect_input_columns(size_type base_column_size,
@@ -274,25 +279,58 @@ void launch_column_output_kernel(jitify2::ConfiguredKernel& kernel,
 template <typename T>
 void launch_span_kernel(jitify2::ConfiguredKernel& kernel,
                         device_span<T> output,
+                        rmm::device_buffer& null_mask,
                         std::vector<column_view> const& input_cols,
                         std::optional<void*> user_data,
                         rmm::cuda_stream_view stream,
                         rmm::device_async_resource_ref mr)
 {
-  auto outputs = to_device_vector(
-    std::vector{cudf::jit::device_span<T>{output.data(), output.size()}}, stream, mr);
+  auto outputs = to_device_vector(std::vector{cudf::jit::device_optional_span<T>{
+                                    cudf::jit::device_span<T>{output.data(), output.size()},
+                                    static_cast<bitmask_type*>(null_mask.data())}},
+                                  stream,
+                                  mr);
+
   auto [input_handles, inputs] =
     column_views_to_device<column_device_view, column_view>(input_cols, stream, mr);
 
-  cudf::jit::device_span<T> const* outputs_ptr = outputs.data();
-  column_device_view const* inputs_ptr         = inputs.data();
-  void* p_user_data                            = user_data.value_or(nullptr);
+  cudf::jit::device_optional_span<T> const* outputs_ptr = outputs.data();
+  column_device_view const* inputs_ptr                  = inputs.data();
+  void* p_user_data                                     = user_data.value_or(nullptr);
 
   std::array<void*, 3> args{&outputs_ptr, &inputs_ptr, &p_user_data};
 
   kernel->launch(args.data());
 }
 
+std::tuple<rmm::device_buffer, size_type> make_transform_null_mask(
+  column_view base_column,
+  std::vector<column_view> const& inputs,
+  rmm::cuda_stream_view stream,
+  rmm::device_async_resource_ref mr)
+{
+  // collect the non-scalar elements that contribute to the resulting bitmask
+  std::vector<column_view> bitmask_columns;
+
+  // to handle null masks for scalars, we just check if the scalar element is null. If it is null,
+  // then all the rows of the transform output will be null. This helps us prevent creating
+  // column-sized bitmasks for each scalar.
+  for (column_view const& col : inputs) {
+    if (is_scalar(base_column.size(), col.size())) {
+      // all nulls
+      if (col.has_nulls()) {
+        return std::make_tuple(
+          create_null_mask(base_column.size(), mask_state::ALL_NULL, stream, mr),
+          base_column.size());
+      }
+    } else {
+      bitmask_columns.emplace_back(col);
+    }
+  }
+
+  return cudf::bitmask_and(table_view{bitmask_columns}, stream, mr);
+}
+
 std::unique_ptr<column> transform_operation(column_view base_column,
                                             data_type output_type,
                                             std::vector<column_view> const& inputs,
@@ -302,14 +340,10 @@ std::unique_ptr<column> transform_operation(column_view base_column,
                                             rmm::cuda_stream_view stream,
                                             rmm::device_async_resource_ref mr)
 {
-  auto output = make_fixed_width_column(output_type,
-                                        base_column.size(),
-                                        copy_bitmask(base_column, stream, mr),
-                                        base_column.null_count(),
-                                        stream,
-                                        mr);
+  auto [null_mask, null_count] = make_transform_null_mask(base_column, inputs, stream, mr);
 
-  if (base_column.is_empty()) { return output; }
+  auto output = make_fixed_width_column(
+    output_type, base_column.size(), std::move(null_mask), null_count, stream, mr);
 
   auto kernel = build_transform_kernel(is_fixed_point(output_type)
                                          ? "cudf::transformation::jit::fixed_point_kernel"
@@ -336,7 +370,7 @@ std::unique_ptr<column> string_view_operation(column_view base_column,
                                               rmm::cuda_stream_view stream,
                                               rmm::device_async_resource_ref mr)
 {
-  if (base_column.is_empty()) { return make_empty_column(data_type{type_id::STRING}); }
+  auto [null_mask, null_count] = make_transform_null_mask(base_column, inputs, stream, mr);
 
   auto kernel = build_span_kernel("cudf::transformation::jit::span_kernel",
                                   base_column.size(),
@@ -350,11 +384,11 @@ std::unique_ptr<column> string_view_operation(column_view base_column,
 
   rmm::device_uvector<string_view> string_views(base_column.size(), stream, mr);
 
-  launch_span_kernel<string_view>(kernel, string_views, inputs, user_data, stream, mr);
+  launch_span_kernel<string_view>(kernel, string_views, null_mask, inputs, user_data, stream, mr);
 
   auto output = make_strings_column(string_views, string_view{}, stream, mr);
 
-  output->set_null_mask(copy_bitmask(base_column, stream, mr), base_column.null_count());
+  output->set_null_mask(std::move(null_mask), null_count);
 
   return output;
 }
@@ -383,12 +417,6 @@ void perform_checks(column_view base_column,
                            }),
                "All transform input columns must have the same size or be scalar (have size 1)",
                std::invalid_argument);
-
-  CUDF_EXPECTS(
-    std::all_of(
-      inputs.begin(), inputs.end(), [&](auto const& input) { return input.null_count() == 0; }),
-    "All transform input columns must have no nulls",
-    std::invalid_argument);
 }
 
 }  // namespace
diff --git a/cpp/tests/transform/integration/unary_transform_test.cpp b/cpp/tests/transform/integration/unary_transform_test.cpp
index cd8a3e2372b..6af098f00eb 100644
--- a/cpp/tests/transform/integration/unary_transform_test.cpp
+++ b/cpp/tests/transform/integration/unary_transform_test.cpp
@@ -100,10 +100,8 @@ TEST_F(AssertsTest, UnequalRowCount)
 
 TEST_F(AssertsTest, NullSupport)
 {
-  EXPECT_THROW(
-    cudf::transform(
-      {a, b_nulls, t}, udf, cudf::data_type{cudf::type_id::FLOAT32}, false, std::nullopt),
-    std::invalid_argument);
+  EXPECT_NO_THROW(cudf::transform(
+    {a, b_nulls, t}, udf, cudf::data_type{cudf::type_id::FLOAT32}, false, std::nullopt));
 }
 
 struct UnaryOperationIntegrationTest : public cudf::test::BaseFixture {
@@ -190,6 +188,9 @@ __device__ inline void    fdsf   (
 
   test_udf<dtype>(cuda.c_str(), op, data_init, 500, false);
   test_udf<dtype>(ptx.c_str(), op, data_init, 500, true);
+
+  test_udf<dtype>(cuda.c_str(), op, data_init, 0, false);
+  test_udf<dtype>(ptx.c_str(), op, data_init, 0, true);
 }
 
 TEST_F(UnaryOperationIntegrationTest, Transform_INT32_INT32)
@@ -574,6 +575,7 @@ TEST_F(StringOperationTest, Output)
   auto b = cudf::test::strings_column_wrapper{"aa", "is", "dd", "ddd", "e", "fff"};
   auto c = cudf::test::strings_column_wrapper{"a", "b", "the", "dddd", "e", "fff"};
   auto d = cudf::test::strings_column_wrapper{"a", "b", "d", "largest", "lexicographical", "test"};
+  auto empty = cudf::test::strings_column_wrapper{};
 
   std::string cuda = R"***(
     __device__ void transform(cudf::string_view * out, cudf::string_view a, cudf::string_view b, cudf::string_view c, cudf::string_view d){
@@ -586,6 +588,12 @@ TEST_F(StringOperationTest, Output)
   auto result = cudf::transform({a, b, c, d}, cuda, cudf::data_type(cudf::type_id::STRING), false);
 
   CUDF_TEST_EXPECT_COLUMNS_EQUAL(expected, result->view());
+
+  auto expected_empty = cudf::test::strings_column_wrapper{};
+  auto result_empty   = cudf::transform(
+    {empty, empty, empty, empty}, cuda, cudf::data_type(cudf::type_id::STRING), false);
+
+  CUDF_TEST_EXPECT_COLUMNS_EQUAL(expected_empty, result_empty->view());
 }
 
 TEST_F(StringOperationTest, StringConcat)
@@ -640,4 +648,197 @@ __device__ void transform(void* user_data, cudf::size_type row,
   CUDF_TEST_EXPECT_COLUMNS_EQUAL(expected, result->view());
 }
 
+struct NullTest : public cudf::test::BaseFixture {
+ protected:
+  void SetUp() override
+  {
+    if (!cudf::is_runtime_jit_supported()) {
+      GTEST_SKIP() << "Skipping tests that require runtime JIT support";
+    }
+  }
+
+  char const* const cuda =
+    "__device__ inline void lerp(float* output, float low, float high, float t){*output = low - t "
+    "* low + t * high; }";
+
+  // Generated from NUMBA, using:
+  //
+  // ```py
+  //
+  // from numba import cuda, float32
+  // from numba.cuda import compile_ptx_for_current_device
+  //
+  // # Define a CUDA device function
+  //
+  // @cuda.jit(device=True)
+  // def op(low, high, t):
+  //         return low - t * low + t * high
+  //
+  // # Define argument types for the function
+  // arg_types = (float32, float32, float32)
+  //
+  // # Compile the device function as relocatable
+  // ptx, _ = cuda.compile_ptx_for_current_device(op, arg_types, device=True)
+  //
+  //
+  // # Print the PTX code
+  // print("Relocatable PTX Code:")
+  // print(ptx)
+  //
+  //
+  // ```
+  //
+  char const* const ptx =
+    R"***(
+//
+// Generated by NVIDIA NVVM Compiler
+//
+// Compiler Build ID: CL-35813241
+// Cuda compilation tools, release 12.9, V12.9.41
+// Based on NVVM 7.0.1
+//
+
+.version 8.8
+.target sm_86
+.address_size 64
+
+// .globl	_ZN8__main__2opB2v1B96cw51cXTLSUwv1sCUt9Ww0FEw09RRQPKiLTj0gIGIFp_2b2oLQFEYYkHSQB1OQAk0Bynm21OizQ1K0UoIGvDpQE8oxrNQE_3dEfff
+.common .global .align 8 .u64 _ZN08NumbaEnv8__main__2opB2v1B96cw51cXTLSUwv1sCUt9Ww0FEw09RRQPKiLTj0gIGIFp_2b2oLQFEYYkHSQB1OQAk0Bynm21OizQ1K0UoIGvDpQE8oxrNQE_3dEfff;
+
+.visible .func  (.param .b32 func_retval0) _ZN8__main__2opB2v1B96cw51cXTLSUwv1sCUt9Ww0FEw09RRQPKiLTj0gIGIFp_2b2oLQFEYYkHSQB1OQAk0Bynm21OizQ1K0UoIGvDpQE8oxrNQE_3dEfff(
+.param .b64 _ZN8__main__2opB2v1B96cw51cXTLSUwv1sCUt9Ww0FEw09RRQPKiLTj0gIGIFp_2b2oLQFEYYkHSQB1OQAk0Bynm21OizQ1K0UoIGvDpQE8oxrNQE_3dEfff_param_0,
+.param .b32 _ZN8__main__2opB2v1B96cw51cXTLSUwv1sCUt9Ww0FEw09RRQPKiLTj0gIGIFp_2b2oLQFEYYkHSQB1OQAk0Bynm21OizQ1K0UoIGvDpQE8oxrNQE_3dEfff_param_1,
+.param .b32 _ZN8__main__2opB2v1B96cw51cXTLSUwv1sCUt9Ww0FEw09RRQPKiLTj0gIGIFp_2b2oLQFEYYkHSQB1OQAk0Bynm21OizQ1K0UoIGvDpQE8oxrNQE_3dEfff_param_2,
+.param .b32 _ZN8__main__2opB2v1B96cw51cXTLSUwv1sCUt9Ww0FEw09RRQPKiLTj0gIGIFp_2b2oLQFEYYkHSQB1OQAk0Bynm21OizQ1K0UoIGvDpQE8oxrNQE_3dEfff_param_3
+)
+{
+.reg .f32 	%f<7>;
+.reg .b32 	%r<2>;
+.reg .b64 	%rd<2>;
+
+
+ld.param.u64 	%rd1, [_ZN8__main__2opB2v1B96cw51cXTLSUwv1sCUt9Ww0FEw09RRQPKiLTj0gIGIFp_2b2oLQFEYYkHSQB1OQAk0Bynm21OizQ1K0UoIGvDpQE8oxrNQE_3dEfff_param_0];
+ld.param.f32 	%f1, [_ZN8__main__2opB2v1B96cw51cXTLSUwv1sCUt9Ww0FEw09RRQPKiLTj0gIGIFp_2b2oLQFEYYkHSQB1OQAk0Bynm21OizQ1K0UoIGvDpQE8oxrNQE_3dEfff_param_1];
+ld.param.f32 	%f2, [_ZN8__main__2opB2v1B96cw51cXTLSUwv1sCUt9Ww0FEw09RRQPKiLTj0gIGIFp_2b2oLQFEYYkHSQB1OQAk0Bynm21OizQ1K0UoIGvDpQE8oxrNQE_3dEfff_param_2];
+ld.param.f32 	%f3, [_ZN8__main__2opB2v1B96cw51cXTLSUwv1sCUt9Ww0FEw09RRQPKiLTj0gIGIFp_2b2oLQFEYYkHSQB1OQAk0Bynm21OizQ1K0UoIGvDpQE8oxrNQE_3dEfff_param_3];
+mul.f32 	%f4, %f1, %f3;
+sub.f32 	%f5, %f1, %f4;
+fma.rn.f32 	%f6, %f2, %f3, %f5;
+st.f32 	[%rd1], %f6;
+mov.u32 	%r1, 0;
+st.param.b32 	[func_retval0+0], %r1;
+ret;
+}
+)***";
+
+  float const LOW         = 100.0F;
+  float const HIGH        = 200.0F;
+  float const T           = 0.25F;
+  float const RES         = LOW - T * LOW + T * HIGH;
+  cudf::size_type const N = 200;
+
+  std::vector<float> low_host      = std::vector<float>(N, LOW);
+  std::vector<float> high_host     = std::vector<float>(N, HIGH);
+  std::vector<float> t_host        = std::vector<float>(N, T);
+  std::vector<float> t_scalar_host = std::vector<float>(1, T);
+  std::vector<float> expected_host = std::vector<float>(N, RES);
+
+  auto fourth()
+  {
+    return cudf::detail::make_counting_transform_iterator(0, [](auto i) { return (i % 4) == 0; });
+  }
+
+  auto fifth()
+  {
+    return cudf::detail::make_counting_transform_iterator(0, [](auto i) { return (i % 5) == 0; });
+  }
+};
+
+TEST_F(NullTest, ColumnNulls)
+{
+  auto low =
+    cudf::test::fixed_width_column_wrapper<float>(low_host.begin(), low_host.end(), fourth())
+      .release();
+  auto high =
+    cudf::test::fixed_width_column_wrapper<float>(high_host.begin(), high_host.end(), fifth())
+      .release();
+  auto t = cudf::test::fixed_width_column_wrapper<float>(t_host.begin(), t_host.end()).release();
+  auto expected =
+    cudf::test::fixed_width_column_wrapper<float>(expected_host.begin(), expected_host.end())
+      .release();
+
+  auto [expected_mask, expect_null_count] = cudf::bitmask_and(cudf::table_view({*low, *high}));
+
+  expected->set_null_mask(std::move(expected_mask), expect_null_count);
+
+  auto cuda_result =
+    cudf::transform({*low, *high, *t}, cuda, cudf::data_type(cudf::type_id::FLOAT32), false);
+
+  CUDF_TEST_EXPECT_COLUMNS_EQUAL(*cuda_result, *expected);
+
+  auto ptx_result =
+    cudf::transform({*low, *high, *t}, ptx, cudf::data_type(cudf::type_id::FLOAT32), true);
+
+  CUDF_TEST_EXPECT_COLUMNS_EQUAL(*ptx_result, *expected);
+}
+
+TEST_F(NullTest, ColumnNulls_And_Scalar)
+{
+  auto low =
+    cudf::test::fixed_width_column_wrapper<float>(low_host.begin(), low_host.end(), fourth())
+      .release();
+  auto high =
+    cudf::test::fixed_width_column_wrapper<float>(high_host.begin(), high_host.end(), fifth())
+      .release();
+  auto t_scalar =
+    cudf::test::fixed_width_column_wrapper<float>(t_scalar_host.begin(), t_scalar_host.end())
+      .release();
+  auto expected =
+    cudf::test::fixed_width_column_wrapper<float>(expected_host.begin(), expected_host.end())
+      .release();
+
+  auto [expected_mask, expect_null_count] = cudf::bitmask_and(cudf::table_view({*low, *high}));
+
+  expected->set_null_mask(std::move(expected_mask), expect_null_count);
+
+  auto cuda_result =
+    cudf::transform({*low, *high, *t_scalar}, cuda, cudf::data_type(cudf::type_id::FLOAT32), false);
+
+  CUDF_TEST_EXPECT_COLUMNS_EQUAL(*cuda_result, *expected);
+
+  auto ptx_result =
+    cudf::transform({*low, *high, *t_scalar}, ptx, cudf::data_type(cudf::type_id::FLOAT32), true);
+
+  CUDF_TEST_EXPECT_COLUMNS_EQUAL(*ptx_result, *expected);
+}
+
+TEST_F(NullTest, ColumnNulls_And_ScalarNull)
+{
+  auto low =
+    cudf::test::fixed_width_column_wrapper<float>(low_host.begin(), low_host.end(), fourth())
+      .release();
+  auto high =
+    cudf::test::fixed_width_column_wrapper<float>(high_host.begin(), high_host.end(), fifth())
+      .release();
+  auto t_scalar = cudf::test::fixed_width_column_wrapper<float>(
+                    t_scalar_host.begin(), t_scalar_host.end(), {false})
+                    .release();
+  auto expected =
+    cudf::test::fixed_width_column_wrapper<float>(expected_host.begin(), expected_host.end())
+      .release();
+
+  expected->set_null_mask(cudf::create_null_mask(low_host.size(), cudf::mask_state::ALL_NULL),
+                          low->size());
+
+  auto cuda_result =
+    cudf::transform({*low, *high, *t_scalar}, cuda, cudf::data_type(cudf::type_id::FLOAT32), false);
+
+  CUDF_TEST_EXPECT_COLUMNS_EQUAL(*cuda_result, *expected);
+
+  auto ptx_result =
+    cudf::transform({*low, *high, *t_scalar}, ptx, cudf::data_type(cudf::type_id::FLOAT32), true);
+
+  CUDF_TEST_EXPECT_COLUMNS_EQUAL(*ptx_result, *expected);
+}
+
 }  // namespace transformation
diff --git a/java/src/test/java/ai/rapids/cudf/ColumnVectorTest.java b/java/src/test/java/ai/rapids/cudf/ColumnVectorTest.java
index 6eb2420f3c4..18393ca7b83 100644
--- a/java/src/test/java/ai/rapids/cudf/ColumnVectorTest.java
+++ b/java/src/test/java/ai/rapids/cudf/ColumnVectorTest.java
@@ -89,10 +89,10 @@ public class ColumnVectorTest extends CudfTestBase {
 
   @Test
   void testTransformVector() {
-    try (ColumnVector cv = ColumnVector.fromBoxedInts(2,3,4);
+    try (ColumnVector cv = ColumnVector.fromBoxedInts(2, 3, null, 4);
          ColumnVector cv1 = cv.transform(ptx, true);
          ColumnVector cv2 = cv.transform(cuda, false);
-         ColumnVector expected = ColumnVector.fromBoxedInts(2*2-2, 3*3-3, 4*4-4)) {
+         ColumnVector expected = ColumnVector.fromBoxedInts(2 * 2 - 2, 3 * 3 - 3, null, 4 * 4 - 4)) {
       assertColumnsAreEqual(expected, cv1);
       assertColumnsAreEqual(expected, cv2);
     }


Reviews:
user: login: davidwendt
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-05-16T11:49:25Z
user: login: lamarrr
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-05-16T14:35:40Z
user: login: davidwendt
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-05-19T12:28:46Z
user: login: lamarrr
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-05-19T12:41:16Z
user: login: ttnghia
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-05-20T17:11:24Z
user: login: ttnghia
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-05-20T17:13:36Z
user: login: ttnghia
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-05-20T17:15:09Z
user: login: lamarrr
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-05-22T14:49:43Z
user: login: lamarrr
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-05-22T14:49:53Z
user: login: ttnghia
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-05-22T17:18:38Z
user: login: davidwendt
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-05-22T17:31:22Z
user: login: ttnghia
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-05-22T23:24:57Z
user: login: lamarrr
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-05-22T23:51:16Z
user: login: davidwendt
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-05-28T17:14:52Z
user: login: kingcrimsontianyu
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-05-28T19:38:04Z
user: login: kingcrimsontianyu
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-05-28T19:43:58Z
user: login: kingcrimsontianyu
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-05-28T19:55:12Z
user: login: kingcrimsontianyu
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-05-28T20:19:02Z
user: login: lamarrr
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-05-29T10:58:59Z
user: login: lamarrr
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-05-29T11:03:58Z
user: login: lamarrr
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-05-29T11:04:06Z
user: login: lamarrr
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-05-29T11:04:39Z
user: login: lamarrr
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-05-29T11:05:02Z
user: login: lamarrr
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-05-29T11:25:15Z
user: login: davidwendt
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-05-29T12:47:02Z
user: login: lamarrr
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-05-29T13:39:52Z
user: login: davidwendt
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-05-29T13:50:16Z
user: login: lamarrr
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-05-29T14:36:00Z
user: login: davidwendt
state: APPROVED
author_association: CONTRIBUTOR
submitted_at: 2025-05-29T16:08:02Z
user: login: kingcrimsontianyu
state: APPROVED
author_association: CONTRIBUTOR
submitted_at: 2025-05-29T17:35:09Z
user: login: mhaseeb123
state: APPROVED
author_association: MEMBER
submitted_at: 2025-05-29T18:48:09Z
user: login: ttnghia
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-05-29T19:43:21Z
user: login: ttnghia
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-05-29T19:45:43Z
user: login: ttnghia
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-05-29T19:46:13Z
user: login: ttnghia
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-05-29T19:46:28Z
user: login: ttnghia
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-05-29T19:47:13Z
user: login: ttnghia
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-05-29T19:48:10Z
user: login: ttnghia
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-05-29T19:48:29Z
user: login: lamarrr
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-05-29T19:54:16Z
user: login: lamarrr
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-05-29T19:57:42Z
user: login: lamarrr
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-05-29T19:58:11Z
user: login: lamarrr
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-05-29T19:59:37Z
user: login: lamarrr
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-05-29T20:07:53Z
user: login: ttnghia
state: APPROVED
author_association: CONTRIBUTOR
submitted_at: 2025-05-29T20:12:00Z

Review comments:
diff_hunk: @@ -121,6 +165,15 @@ CUDF_KERNEL void kernel(cudf::mutable_column_device_view_core const* outputs,
   thread_index_type const size   = outputs[0].size();
 
   for (auto i = start; i < size; i += stride) {
+    if constexpr (has_nulls) {
+      if ((false || ... || In::is_null(inputs, i))) {
+        Out::set_null(outputs, i);
+        continue;
+      } else {
+        Out::set_valid(outputs, i);
+      }
+    }
path: cpp/src/transform/jit/kernel.cu
user: login: davidwendt
body: I think this would be less efficient than just gathering the bitmasks separately and applying it to the output column. Also, I think this will be a race condition since multiple threads may be trying to update the same bytes.
created_at: 2025-05-16T11:48:51Z
author_association: CONTRIBUTOR
diff_hunk: @@ -121,6 +165,15 @@ CUDF_KERNEL void kernel(cudf::mutable_column_device_view_core const* outputs,
   thread_index_type const size   = outputs[0].size();
 
   for (auto i = start; i < size; i += stride) {
+    if constexpr (has_nulls) {
+      if ((false || ... || In::is_null(inputs, i))) {
+        Out::set_null(outputs, i);
+        continue;
+      } else {
+        Out::set_valid(outputs, i);
+      }
+    }
path: cpp/src/transform/jit/kernel.cu
user: login: lamarrr
body: > I think this would be less efficient than just gathering the bitmasks separately and applying it to the output column

Ok, I'll do that.

> Also, I think this will be a race condition since multiple threads may be trying to update the same bytes.

It's an atomic bit operation
created_at: 2025-05-16T14:35:40Z
author_association: CONTRIBUTOR
diff_hunk: @@ -121,6 +165,15 @@ CUDF_KERNEL void kernel(cudf::mutable_column_device_view_core const* outputs,
   thread_index_type const size   = outputs[0].size();
 
   for (auto i = start; i < size; i += stride) {
+    if constexpr (has_nulls) {
+      if ((false || ... || In::is_null(inputs, i))) {
+        Out::set_null(outputs, i);
+        continue;
+      } else {
+        Out::set_valid(outputs, i);
+      }
+    }
path: cpp/src/transform/jit/kernel.cu
user: login: davidwendt
body: Also, we have found there is no significant performance improvement for making `has_nulls` a template argument where the code divergence is very short -- like in this case where the bit is just checked. I'd rather the `has_nulls` be a runtime parameter or just check for the validity-mask all the time. It is not worth (compiler) generating separating kernels for just the single-bit check.
Perhaps some benchmarks that prove this worthy of a template parameter?
created_at: 2025-05-19T12:28:45Z
author_association: CONTRIBUTOR
diff_hunk: @@ -121,6 +165,15 @@ CUDF_KERNEL void kernel(cudf::mutable_column_device_view_core const* outputs,
   thread_index_type const size   = outputs[0].size();
 
   for (auto i = start; i < size; i += stride) {
+    if constexpr (has_nulls) {
+      if ((false || ... || In::is_null(inputs, i))) {
+        Out::set_null(outputs, i);
+        continue;
+      } else {
+        Out::set_valid(outputs, i);
+      }
+    }
path: cpp/src/transform/jit/kernel.cu
user: login: lamarrr
body: > Also, we have found there is no significant performance improvement for making has_nulls a template argument where the code divergence is very short -- like in this case where the bit is just checked.

Got it!
I had the impression it may be worth it since that was how AST was also handling nulls. I don't think it will have much of an impact either.
created_at: 2025-05-19T12:41:16Z
author_association: CONTRIBUTOR
diff_hunk: @@ -97,5 +98,52 @@ struct device_span {
   CUDF_HOST_DEVICE [[nodiscard]] constexpr element_type* end() const { return _data + _size; }
 };
 
+/**
+ * @brief A span type with optional elements.
+ *
+ */
path: cpp/include/cudf/jit/span.cuh
user: login: ttnghia
body: ```suggestion
/**
 * @brief A span type with optional elements.
 */
```
created_at: 2025-05-20T17:11:24Z
author_association: CONTRIBUTOR
diff_hunk: @@ -97,5 +98,52 @@ struct device_span {
   CUDF_HOST_DEVICE [[nodiscard]] constexpr element_type* end() const { return _data + _size; }
 };
 
+/**
+ * @brief A span type with optional elements.
+ *
+ */
path: cpp/include/cudf/jit/span.cuh
user: login: ttnghia
body: What does optional mean here? I don't see any optional.
created_at: 2025-05-20T17:13:36Z
author_association: CONTRIBUTOR
diff_hunk: @@ -86,7 +99,7 @@ struct scalar {
   static constexpr int32_t index = Accessor::index;
 
   static __device__ decltype(auto) element(cudf::mutable_column_device_view_core const* outputs,
-                                           cudf::size_type row)
+                                           [[maybe_unused]] cudf::size_type row)
path: cpp/src/transform/jit/kernel.cu
user: login: ttnghia
body: If the parameter is unused, do not mention it instead.
```suggestion
                                           cudf::size_type)
```
created_at: 2025-05-20T17:15:08Z
author_association: CONTRIBUTOR
diff_hunk: @@ -97,5 +98,52 @@ struct device_span {
   CUDF_HOST_DEVICE [[nodiscard]] constexpr element_type* end() const { return _data + _size; }
 };
 
+/**
+ * @brief A span type with optional elements.
+ *
+ */
path: cpp/include/cudf/jit/span.cuh
user: login: lamarrr
body: optional means the same as before; it contains a value or a null. In this case, it is an array/span of them.
created_at: 2025-05-22T14:49:43Z
author_association: CONTRIBUTOR
diff_hunk: @@ -86,7 +99,7 @@ struct scalar {
   static constexpr int32_t index = Accessor::index;
 
   static __device__ decltype(auto) element(cudf::mutable_column_device_view_core const* outputs,
-                                           cudf::size_type row)
+                                           [[maybe_unused]] cudf::size_type row)
path: cpp/src/transform/jit/kernel.cu
user: login: lamarrr
body: ok
created_at: 2025-05-22T14:49:53Z
author_association: CONTRIBUTOR
diff_hunk: @@ -97,5 +98,52 @@ struct device_span {
   CUDF_HOST_DEVICE [[nodiscard]] constexpr element_type* end() const { return _data + _size; }
 };
 
+/**
+ * @brief A span type with optional elements.
+ *
+ */
path: cpp/include/cudf/jit/span.cuh
user: login: ttnghia
body: That is not similar to an optional. For null, the array element still exists, and it is a null instead of not available as an optional. I think the better name should be something like `nullable_span` or `span_with_null` or something else.
created_at: 2025-05-22T17:18:37Z
author_association: CONTRIBUTOR
diff_hunk: @@ -97,5 +98,52 @@ struct device_span {
   CUDF_HOST_DEVICE [[nodiscard]] constexpr element_type* end() const { return _data + _size; }
 };
 
+/**
+ * @brief A span type with optional elements.
+ *
+ */
path: cpp/include/cudf/jit/span.cuh
user: login: davidwendt
body: The `optional` is the correct term for libcudf. It is inline with optional-iterators we use for identifying null rows in a column. https://github.com/rapidsai/cudf/blob/c9231de061967575b293de03863c4d55e1e885eb/cpp/include/cudf/column/column_device_view.cuh#L362
created_at: 2025-05-22T17:31:22Z
author_association: CONTRIBUTOR
diff_hunk: @@ -97,5 +98,52 @@ struct device_span {
   CUDF_HOST_DEVICE [[nodiscard]] constexpr element_type* end() const { return _data + _size; }
 };
 
+/**
+ * @brief A span type with optional elements.
+ *
+ */
path: cpp/include/cudf/jit/span.cuh
user: login: ttnghia
body: If want to keep the current name, please add more comment to clarify what it is. Just "optional elements" seems very confusing.
created_at: 2025-05-22T23:24:57Z
author_association: CONTRIBUTOR
diff_hunk: @@ -97,5 +98,52 @@ struct device_span {
   CUDF_HOST_DEVICE [[nodiscard]] constexpr element_type* end() const { return _data + _size; }
 };
 
+/**
+ * @brief A span type with optional elements.
+ *
+ */
path: cpp/include/cudf/jit/span.cuh
user: login: lamarrr
body: done. 


Yeah, I understand. I suppose that's because of its internal representation. Most would assume it is a `span<std::optional>`, it's just like `std::vector<bool>`
created_at: 2025-05-22T23:51:16Z
author_association: CONTRIBUTOR
diff_hunk: @@ -97,5 +98,55 @@ struct device_span {
   CUDF_HOST_DEVICE [[nodiscard]] constexpr element_type* end() const { return _data + _size; }
 };
 
+/**
path: cpp/include/cudf/jit/span.cuh
user: login: davidwendt
body: Does this header file need to be public? I don't think we expect callers to use it. It is just internally by the transform.cpp implementation. I don't think we need the `CUDF_EXPORT` either.
created_at: 2025-05-28T17:14:51Z
author_association: CONTRIBUTOR
diff_hunk: @@ -97,5 +98,55 @@ struct device_span {
   CUDF_HOST_DEVICE [[nodiscard]] constexpr element_type* end() const { return _data + _size; }
 };
 
+/**
+ * @brief A span type with optional/nullable elements.
+ * Optional implies the span contains nullable elements.
+ * The nullability of the elements is internally represented by an optional bitmask which is nullptr
+ * when all the elements are empty.
path: cpp/include/cudf/jit/span.cuh
user: login: kingcrimsontianyu
body: Should it be "when all the elements are valid"?
created_at: 2025-05-28T19:38:04Z
author_association: CONTRIBUTOR
diff_hunk: @@ -97,5 +98,55 @@ struct device_span {
   CUDF_HOST_DEVICE [[nodiscard]] constexpr element_type* end() const { return _data + _size; }
 };
 
+/**
+ * @brief A span type with optional/nullable elements.
+ * Optional implies the span contains nullable elements.
+ * The nullability of the elements is internally represented by an optional bitmask which is nullptr
+ * when all the elements are empty.
+ *
+ */
+template <typename T>
+struct device_optional_span : device_span<T> {
+ private:
+  using base               = device_span<T>;
+  bitmask_type* _null_mask = nullptr;
+
+ public:
+  CUDF_HOST_DEVICE constexpr device_optional_span() {}
+
+  /**
+   * @brief Constructs an optional span from a span and a null-mask.
+   *
+   * @param span Span containing the elements
+   * @param null_mask The null-mask determining the validity of the elements or nullptr if all
+   * valid.
+   */
+  CUDF_HOST_DEVICE device_optional_span(device_span<T> span, bitmask_type* null_mask)
+    : base{span}, _null_mask{null_mask}
+  {
+  }
+
+  /// @copydoc column_device_view::nullable
+  [[nodiscard]] CUDF_HOST_DEVICE bool nullable() const { return _null_mask != nullptr; }
+
+  /// @copydoc column_device_view::is_valid_nocheck
+  [[nodiscard]] __device__ bool is_valid_nocheck(size_t element_index) const
+  {
+    return bit_is_set(_null_mask, element_index);
+  }
+
+  /// @copydoc column_device_view::is_valid
+  [[nodiscard]] __device__ bool is_valid(size_t element_index) const
+  {
+    return not nullable() or is_valid_nocheck(element_index);
+  }
+
+  /// @copydoc column_device_view::is_null
+  [[nodiscard]] __device__ bool is_null(size_t element_index) const
+  {
+    return !is_valid(element_index);
path: cpp/include/cudf/jit/span.cuh
user: login: kingcrimsontianyu
body: Nit: `not is_valid` per cudf convention.
created_at: 2025-05-28T19:43:58Z
author_association: CONTRIBUTOR
diff_hunk: @@ -274,25 +274,63 @@ void launch_column_output_kernel(jitify2::ConfiguredKernel& kernel,
 template <typename T>
 void launch_span_kernel(jitify2::ConfiguredKernel& kernel,
                         device_span<T> output,
+                        rmm::device_buffer& null_mask,
                         std::vector<column_view> const& input_cols,
                         std::optional<void*> user_data,
                         rmm::cuda_stream_view stream,
                         rmm::device_async_resource_ref mr)
 {
-  auto outputs = to_device_vector(
-    std::vector{cudf::jit::device_span<T>{output.data(), output.size()}}, stream, mr);
+  auto outputs = to_device_vector(std::vector{cudf::jit::device_optional_span<T>{
+                                    cudf::jit::device_span<T>{output.data(), output.size()},
+                                    static_cast<bitmask_type*>(null_mask.data())}},
+                                  stream,
+                                  mr);
+
   auto [input_handles, inputs] =
     column_views_to_device<column_device_view, column_view>(input_cols, stream, mr);
 
-  cudf::jit::device_span<T> const* outputs_ptr = outputs.data();
-  column_device_view const* inputs_ptr         = inputs.data();
-  void* p_user_data                            = user_data.value_or(nullptr);
+  cudf::jit::device_optional_span<T> const* outputs_ptr = outputs.data();
+  column_device_view const* inputs_ptr                  = inputs.data();
+  void* p_user_data                                     = user_data.value_or(nullptr);
 
   std::array<void*, 3> args{&outputs_ptr, &inputs_ptr, &p_user_data};
 
   kernel->launch(args.data());
 }
 
+bool is_scalar(cudf::size_type base_column_size, cudf::size_type column_size)
+{
+  return column_size == 1 && column_size != base_column_size;
path: cpp/src/transform/transform.cpp
user: login: kingcrimsontianyu
body: Same nit: `and`
created_at: 2025-05-28T19:55:12Z
author_association: CONTRIBUTOR
diff_hunk: @@ -274,25 +274,63 @@ void launch_column_output_kernel(jitify2::ConfiguredKernel& kernel,
 template <typename T>
 void launch_span_kernel(jitify2::ConfiguredKernel& kernel,
                         device_span<T> output,
+                        rmm::device_buffer& null_mask,
                         std::vector<column_view> const& input_cols,
                         std::optional<void*> user_data,
                         rmm::cuda_stream_view stream,
                         rmm::device_async_resource_ref mr)
 {
-  auto outputs = to_device_vector(
-    std::vector{cudf::jit::device_span<T>{output.data(), output.size()}}, stream, mr);
+  auto outputs = to_device_vector(std::vector{cudf::jit::device_optional_span<T>{
+                                    cudf::jit::device_span<T>{output.data(), output.size()},
+                                    static_cast<bitmask_type*>(null_mask.data())}},
+                                  stream,
+                                  mr);
+
   auto [input_handles, inputs] =
     column_views_to_device<column_device_view, column_view>(input_cols, stream, mr);
 
-  cudf::jit::device_span<T> const* outputs_ptr = outputs.data();
-  column_device_view const* inputs_ptr         = inputs.data();
-  void* p_user_data                            = user_data.value_or(nullptr);
+  cudf::jit::device_optional_span<T> const* outputs_ptr = outputs.data();
+  column_device_view const* inputs_ptr                  = inputs.data();
+  void* p_user_data                                     = user_data.value_or(nullptr);
 
   std::array<void*, 3> args{&outputs_ptr, &inputs_ptr, &p_user_data};
 
   kernel->launch(args.data());
 }
 
+bool is_scalar(cudf::size_type base_column_size, cudf::size_type column_size)
+{
+  return column_size == 1 && column_size != base_column_size;
+}
+
+std::tuple<rmm::device_buffer, size_type> make_transform_null_mask(
+  column_view base_column,
+  std::vector<column_view> const& inputs,
+  rmm::cuda_stream_view stream,
+  rmm::device_async_resource_ref mr)
+{
+  // collect the non-scalar elements that contribute to the resulting bitmask
+  std::vector<column_view> bitmask_columns;
+
+  // to handle null masks for scalars, we just check if the scalar element is null. If it is null,
+  // then all the rows of the transform output will be null. This helps us prevent creating
+  // column-sized bitmasks for each scalar.
+  for (column_view const& col : inputs) {
+    if (is_scalar(base_column.size(), col.size())) {
+      // all nulls
+      if (col.has_nulls()) {
+        return std::make_tuple(
+          create_null_mask(base_column.size(), mask_state::ALL_NULL, stream, mr),
+          base_column.size());
path: cpp/src/transform/transform.cpp
user: login: kingcrimsontianyu
body: ```suggestion
        return {create_null_mask(base_column.size(), mask_state::ALL_NULL, stream, mr),
                base_column.size()};
```
created_at: 2025-05-28T20:19:02Z
author_association: CONTRIBUTOR
diff_hunk: @@ -274,25 +274,63 @@ void launch_column_output_kernel(jitify2::ConfiguredKernel& kernel,
 template <typename T>
 void launch_span_kernel(jitify2::ConfiguredKernel& kernel,
                         device_span<T> output,
+                        rmm::device_buffer& null_mask,
                         std::vector<column_view> const& input_cols,
                         std::optional<void*> user_data,
                         rmm::cuda_stream_view stream,
                         rmm::device_async_resource_ref mr)
 {
-  auto outputs = to_device_vector(
-    std::vector{cudf::jit::device_span<T>{output.data(), output.size()}}, stream, mr);
+  auto outputs = to_device_vector(std::vector{cudf::jit::device_optional_span<T>{
+                                    cudf::jit::device_span<T>{output.data(), output.size()},
+                                    static_cast<bitmask_type*>(null_mask.data())}},
+                                  stream,
+                                  mr);
+
   auto [input_handles, inputs] =
     column_views_to_device<column_device_view, column_view>(input_cols, stream, mr);
 
-  cudf::jit::device_span<T> const* outputs_ptr = outputs.data();
-  column_device_view const* inputs_ptr         = inputs.data();
-  void* p_user_data                            = user_data.value_or(nullptr);
+  cudf::jit::device_optional_span<T> const* outputs_ptr = outputs.data();
+  column_device_view const* inputs_ptr                  = inputs.data();
+  void* p_user_data                                     = user_data.value_or(nullptr);
 
   std::array<void*, 3> args{&outputs_ptr, &inputs_ptr, &p_user_data};
 
   kernel->launch(args.data());
 }
 
+bool is_scalar(cudf::size_type base_column_size, cudf::size_type column_size)
+{
+  return column_size == 1 && column_size != base_column_size;
+}
+
+std::tuple<rmm::device_buffer, size_type> make_transform_null_mask(
+  column_view base_column,
+  std::vector<column_view> const& inputs,
+  rmm::cuda_stream_view stream,
+  rmm::device_async_resource_ref mr)
+{
+  // collect the non-scalar elements that contribute to the resulting bitmask
+  std::vector<column_view> bitmask_columns;
+
+  // to handle null masks for scalars, we just check if the scalar element is null. If it is null,
+  // then all the rows of the transform output will be null. This helps us prevent creating
+  // column-sized bitmasks for each scalar.
+  for (column_view const& col : inputs) {
+    if (is_scalar(base_column.size(), col.size())) {
+      // all nulls
+      if (col.has_nulls()) {
+        return std::make_tuple(
+          create_null_mask(base_column.size(), mask_state::ALL_NULL, stream, mr),
+          base_column.size());
path: cpp/src/transform/transform.cpp
user: login: lamarrr
body: nit?
Depending on the context, brace-initializers without the types specified can be ambiguous; std::make_tuple preserves the value categories and ensures correct construction of the types
created_at: 2025-05-29T10:58:59Z
author_association: CONTRIBUTOR
diff_hunk: @@ -97,5 +98,55 @@ struct device_span {
   CUDF_HOST_DEVICE [[nodiscard]] constexpr element_type* end() const { return _data + _size; }
 };
 
+/**
+ * @brief A span type with optional/nullable elements.
+ * Optional implies the span contains nullable elements.
+ * The nullability of the elements is internally represented by an optional bitmask which is nullptr
+ * when all the elements are empty.
+ *
+ */
+template <typename T>
+struct device_optional_span : device_span<T> {
+ private:
+  using base               = device_span<T>;
+  bitmask_type* _null_mask = nullptr;
+
+ public:
+  CUDF_HOST_DEVICE constexpr device_optional_span() {}
+
+  /**
+   * @brief Constructs an optional span from a span and a null-mask.
+   *
+   * @param span Span containing the elements
+   * @param null_mask The null-mask determining the validity of the elements or nullptr if all
+   * valid.
+   */
+  CUDF_HOST_DEVICE device_optional_span(device_span<T> span, bitmask_type* null_mask)
+    : base{span}, _null_mask{null_mask}
+  {
+  }
+
+  /// @copydoc column_device_view::nullable
+  [[nodiscard]] CUDF_HOST_DEVICE bool nullable() const { return _null_mask != nullptr; }
+
+  /// @copydoc column_device_view::is_valid_nocheck
+  [[nodiscard]] __device__ bool is_valid_nocheck(size_t element_index) const
+  {
+    return bit_is_set(_null_mask, element_index);
+  }
+
+  /// @copydoc column_device_view::is_valid
+  [[nodiscard]] __device__ bool is_valid(size_t element_index) const
+  {
+    return not nullable() or is_valid_nocheck(element_index);
+  }
+
+  /// @copydoc column_device_view::is_null
+  [[nodiscard]] __device__ bool is_null(size_t element_index) const
+  {
+    return !is_valid(element_index);
path: cpp/include/cudf/jit/span.cuh
user: login: lamarrr
body: IIRC `not`  is implemented as macros in some compilers, also depending on the standards version (https://en.wikipedia.org/wiki/C_alternative_tokens).
The linter doesn't pick this up. Is it a documented convention? I don't see how `not` is more readable than `!`.
Or how obvious or readable `comp` is as a replacement for `~`.
created_at: 2025-05-29T11:03:57Z
author_association: CONTRIBUTOR
diff_hunk: @@ -274,25 +274,63 @@ void launch_column_output_kernel(jitify2::ConfiguredKernel& kernel,
 template <typename T>
 void launch_span_kernel(jitify2::ConfiguredKernel& kernel,
                         device_span<T> output,
+                        rmm::device_buffer& null_mask,
                         std::vector<column_view> const& input_cols,
                         std::optional<void*> user_data,
                         rmm::cuda_stream_view stream,
                         rmm::device_async_resource_ref mr)
 {
-  auto outputs = to_device_vector(
-    std::vector{cudf::jit::device_span<T>{output.data(), output.size()}}, stream, mr);
+  auto outputs = to_device_vector(std::vector{cudf::jit::device_optional_span<T>{
+                                    cudf::jit::device_span<T>{output.data(), output.size()},
+                                    static_cast<bitmask_type*>(null_mask.data())}},
+                                  stream,
+                                  mr);
+
   auto [input_handles, inputs] =
     column_views_to_device<column_device_view, column_view>(input_cols, stream, mr);
 
-  cudf::jit::device_span<T> const* outputs_ptr = outputs.data();
-  column_device_view const* inputs_ptr         = inputs.data();
-  void* p_user_data                            = user_data.value_or(nullptr);
+  cudf::jit::device_optional_span<T> const* outputs_ptr = outputs.data();
+  column_device_view const* inputs_ptr                  = inputs.data();
+  void* p_user_data                                     = user_data.value_or(nullptr);
 
   std::array<void*, 3> args{&outputs_ptr, &inputs_ptr, &p_user_data};
 
   kernel->launch(args.data());
 }
 
+bool is_scalar(cudf::size_type base_column_size, cudf::size_type column_size)
+{
+  return column_size == 1 && column_size != base_column_size;
path: cpp/src/transform/transform.cpp
user: login: lamarrr
body: same as above
created_at: 2025-05-29T11:04:06Z
author_association: CONTRIBUTOR
diff_hunk: @@ -97,5 +98,55 @@ struct device_span {
   CUDF_HOST_DEVICE [[nodiscard]] constexpr element_type* end() const { return _data + _size; }
 };
 
+/**
+ * @brief A span type with optional/nullable elements.
+ * Optional implies the span contains nullable elements.
+ * The nullability of the elements is internally represented by an optional bitmask which is nullptr
+ * when all the elements are empty.
path: cpp/include/cudf/jit/span.cuh
user: login: lamarrr
body: yes, thanks!
created_at: 2025-05-29T11:04:39Z
author_association: CONTRIBUTOR
diff_hunk: @@ -97,5 +98,55 @@ struct device_span {
   CUDF_HOST_DEVICE [[nodiscard]] constexpr element_type* end() const { return _data + _size; }
 };
 
+/**
path: cpp/include/cudf/jit/span.cuh
user: login: lamarrr
body: yeah
created_at: 2025-05-29T11:05:02Z
author_association: CONTRIBUTOR
diff_hunk: @@ -97,5 +98,55 @@ struct device_span {
   CUDF_HOST_DEVICE [[nodiscard]] constexpr element_type* end() const { return _data + _size; }
 };
 
+/**
path: cpp/include/cudf/jit/span.cuh
user: login: lamarrr
body: thanks!
I initially did that because it was being used by JITIFY. Seems JITIFY picks it up during preprocessing regardless.
created_at: 2025-05-29T11:25:15Z
author_association: CONTRIBUTOR
diff_hunk: @@ -302,14 +340,10 @@ std::unique_ptr<column> transform_operation(column_view base_column,
                                             rmm::cuda_stream_view stream,
                                             rmm::device_async_resource_ref mr)
 {
-  auto output = make_fixed_width_column(output_type,
-                                        base_column.size(),
-                                        copy_bitmask(base_column, stream, mr),
-                                        base_column.null_count(),
-                                        stream,
-                                        mr);
+  auto [null_mask, null_count] = make_transform_null_mask(base_column, inputs, stream, mr);
 
-  if (base_column.is_empty()) { return output; }
path: cpp/src/transform/transform.cpp
user: login: davidwendt
body: What happens now if the column is empty here? Do we have a test for this?
created_at: 2025-05-29T12:47:02Z
author_association: CONTRIBUTOR
diff_hunk: @@ -302,14 +340,10 @@ std::unique_ptr<column> transform_operation(column_view base_column,
                                             rmm::cuda_stream_view stream,
                                             rmm::device_async_resource_ref mr)
 {
-  auto output = make_fixed_width_column(output_type,
-                                        base_column.size(),
-                                        copy_bitmask(base_column, stream, mr),
-                                        base_column.null_count(),
-                                        stream,
-                                        mr);
+  auto [null_mask, null_count] = make_transform_null_mask(base_column, inputs, stream, mr);
 
-  if (base_column.is_empty()) { return output; }
path: cpp/src/transform/transform.cpp
user: login: lamarrr
body: Presently, it launches the kernel regardless. I don't think having a branch for a zero-sized column is necessary or will yield any practical benefit, especially if it's not a tested codepath.
I've added tests for it now, we aren't performing any operation that should fail for 0-sized arrays/columns.
created_at: 2025-05-29T13:39:52Z
author_association: CONTRIBUTOR
diff_hunk: @@ -302,14 +340,10 @@ std::unique_ptr<column> transform_operation(column_view base_column,
                                             rmm::cuda_stream_view stream,
                                             rmm::device_async_resource_ref mr)
 {
-  auto output = make_fixed_width_column(output_type,
-                                        base_column.size(),
-                                        copy_bitmask(base_column, stream, mr),
-                                        base_column.null_count(),
-                                        stream,
-                                        mr);
+  auto [null_mask, null_count] = make_transform_null_mask(base_column, inputs, stream, mr);
 
-  if (base_column.is_empty()) { return output; }
path: cpp/src/transform/transform.cpp
user: login: davidwendt
body: How many threads are launched? We have had issues where launching a kernel with no threads results in a CUDA error (usually returned in the next CUDA call).
created_at: 2025-05-29T13:50:16Z
author_association: CONTRIBUTOR
diff_hunk: @@ -302,14 +340,10 @@ std::unique_ptr<column> transform_operation(column_view base_column,
                                             rmm::cuda_stream_view stream,
                                             rmm::device_async_resource_ref mr)
 {
-  auto output = make_fixed_width_column(output_type,
-                                        base_column.size(),
-                                        copy_bitmask(base_column, stream, mr),
-                                        base_column.null_count(),
-                                        stream,
-                                        mr);
+  auto [null_mask, null_count] = make_transform_null_mask(base_column, inputs, stream, mr);
 
-  if (base_column.is_empty()) { return output; }
path: cpp/src/transform/transform.cpp
user: login: lamarrr
body: It launches the maximum number of threads needed for maximum occupancy on the GPU (always non-zero). On GH200, for example, that'll be 1024 blockDim and 132 gridDim.
created_at: 2025-05-29T14:36:00Z
author_association: CONTRIBUTOR
diff_hunk: @@ -97,5 +98,59 @@ struct device_span {
   CUDF_HOST_DEVICE [[nodiscard]] constexpr element_type* end() const { return _data + _size; }
 };
 
+/**
+ * @brief A span type with optional/nullable elements.
+ * Optional implies the span contains nullable elements.
+ * The nullability of the elements is internally represented by an optional bitmask which can be
+ * nullptr when all the elements are non-null.
+ *
+ */
path: cpp/src/jit/span.cuh
user: login: ttnghia
body: ```suggestion

/**
 * @brief A span type with optional/nullable elements.
 *
 * Optional implies the span contains nullable elements.
 * The nullability of the elements is internally represented by an optional bitmask which can be
 * nullptr when all the elements are non-null.
 */
```
created_at: 2025-05-29T19:43:21Z
author_association: CONTRIBUTOR
diff_hunk: @@ -274,25 +279,58 @@ void launch_column_output_kernel(jitify2::ConfiguredKernel& kernel,
 template <typename T>
 void launch_span_kernel(jitify2::ConfiguredKernel& kernel,
                         device_span<T> output,
+                        rmm::device_buffer& null_mask,
                         std::vector<column_view> const& input_cols,
                         std::optional<void*> user_data,
                         rmm::cuda_stream_view stream,
                         rmm::device_async_resource_ref mr)
 {
-  auto outputs = to_device_vector(
-    std::vector{cudf::jit::device_span<T>{output.data(), output.size()}}, stream, mr);
+  auto outputs = to_device_vector(std::vector{cudf::jit::device_optional_span<T>{
+                                    cudf::jit::device_span<T>{output.data(), output.size()},
+                                    static_cast<bitmask_type*>(null_mask.data())}},
+                                  stream,
+                                  mr);
+
   auto [input_handles, inputs] =
     column_views_to_device<column_device_view, column_view>(input_cols, stream, mr);
 
-  cudf::jit::device_span<T> const* outputs_ptr = outputs.data();
-  column_device_view const* inputs_ptr         = inputs.data();
-  void* p_user_data                            = user_data.value_or(nullptr);
+  cudf::jit::device_optional_span<T> const* outputs_ptr = outputs.data();
path: cpp/src/transform/transform.cpp
user: login: ttnghia
body: So this is a pointer to a span array?
created_at: 2025-05-29T19:45:43Z
author_association: CONTRIBUTOR
diff_hunk: @@ -302,14 +340,10 @@ std::unique_ptr<column> transform_operation(column_view base_column,
                                             rmm::cuda_stream_view stream,
                                             rmm::device_async_resource_ref mr)
 {
-  auto output = make_fixed_width_column(output_type,
-                                        base_column.size(),
-                                        copy_bitmask(base_column, stream, mr),
-                                        base_column.null_count(),
-                                        stream,
-                                        mr);
+  auto [null_mask, null_count] = make_transform_null_mask(base_column, inputs, stream, mr);
 
-  if (base_column.is_empty()) { return output; }
+  auto output = make_fixed_width_column(
+    output_type, base_column.size(), std::move(null_mask), null_count, stream, mr);
path: cpp/src/transform/transform.cpp
user: login: ttnghia
body: Can we ignore null mask if `null_count == 0`?
created_at: 2025-05-29T19:46:13Z
author_association: CONTRIBUTOR
diff_hunk: @@ -350,11 +384,11 @@ std::unique_ptr<column> string_view_operation(column_view base_column,
 
   rmm::device_uvector<string_view> string_views(base_column.size(), stream, mr);
 
-  launch_span_kernel<string_view>(kernel, string_views, inputs, user_data, stream, mr);
+  launch_span_kernel<string_view>(kernel, string_views, null_mask, inputs, user_data, stream, mr);
 
   auto output = make_strings_column(string_views, string_view{}, stream, mr);
 
-  output->set_null_mask(copy_bitmask(base_column, stream, mr), base_column.null_count());
+  output->set_null_mask(std::move(null_mask), null_count);
path: cpp/src/transform/transform.cpp
user: login: ttnghia
body: Can we ignore this if `null_count == 0`?
```suggestion
if(null_count) { output->set_null_mask(); }
```
created_at: 2025-05-29T19:46:28Z
author_association: CONTRIBUTOR
diff_hunk: @@ -640,4 +648,197 @@ __device__ void transform(void* user_data, cudf::size_type row,
   CUDF_TEST_EXPECT_COLUMNS_EQUAL(expected, result->view());
 }
 
+struct NullTest : public cudf::test::BaseFixture {
+ protected:
+  void SetUp() override
+  {
+    if (!cudf::is_runtime_jit_supported()) {
+      GTEST_SKIP() << "Skipping tests that require runtime JIT support";
+    }
+  }
+
+  char const* const cuda =
+    "__device__ inline void lerp(float* output, float low, float high, float t){*output = low - t "
+    "* low + t * high; }";
+
+  // Generated from NUMBA, using:
+  //
+  // ```py
+  //
+  // from numba import cuda, float32
+  // from numba.cuda import compile_ptx_for_current_device
+  //
+  // # Define a CUDA device function
+  //
+  // @cuda.jit(device=True)
+  // def op(low, high, t):
+  //         return low - t * low + t * high
+  //
+  // # Define argument types for the function
+  // arg_types = (float32, float32, float32)
+  //
+  // # Compile the device function as relocatable
+  // ptx, _ = cuda.compile_ptx_for_current_device(op, arg_types, device=True)
+  //
+  //
+  // # Print the PTX code
+  // print("Relocatable PTX Code:")
+  // print(ptx)
+  //
+  //
+  // ```
+  //
+  char const* const ptx =
+    R"***(
+//
+// Generated by NVIDIA NVVM Compiler
+//
+// Compiler Build ID: CL-35813241
+// Cuda compilation tools, release 12.9, V12.9.41
+// Based on NVVM 7.0.1
+//
+
+.version 8.8
+.target sm_86
+.address_size 64
+
+// .globl	_ZN8__main__2opB2v1B96cw51cXTLSUwv1sCUt9Ww0FEw09RRQPKiLTj0gIGIFp_2b2oLQFEYYkHSQB1OQAk0Bynm21OizQ1K0UoIGvDpQE8oxrNQE_3dEfff
+.common .global .align 8 .u64 _ZN08NumbaEnv8__main__2opB2v1B96cw51cXTLSUwv1sCUt9Ww0FEw09RRQPKiLTj0gIGIFp_2b2oLQFEYYkHSQB1OQAk0Bynm21OizQ1K0UoIGvDpQE8oxrNQE_3dEfff;
+
+.visible .func  (.param .b32 func_retval0) _ZN8__main__2opB2v1B96cw51cXTLSUwv1sCUt9Ww0FEw09RRQPKiLTj0gIGIFp_2b2oLQFEYYkHSQB1OQAk0Bynm21OizQ1K0UoIGvDpQE8oxrNQE_3dEfff(
+.param .b64 _ZN8__main__2opB2v1B96cw51cXTLSUwv1sCUt9Ww0FEw09RRQPKiLTj0gIGIFp_2b2oLQFEYYkHSQB1OQAk0Bynm21OizQ1K0UoIGvDpQE8oxrNQE_3dEfff_param_0,
+.param .b32 _ZN8__main__2opB2v1B96cw51cXTLSUwv1sCUt9Ww0FEw09RRQPKiLTj0gIGIFp_2b2oLQFEYYkHSQB1OQAk0Bynm21OizQ1K0UoIGvDpQE8oxrNQE_3dEfff_param_1,
+.param .b32 _ZN8__main__2opB2v1B96cw51cXTLSUwv1sCUt9Ww0FEw09RRQPKiLTj0gIGIFp_2b2oLQFEYYkHSQB1OQAk0Bynm21OizQ1K0UoIGvDpQE8oxrNQE_3dEfff_param_2,
+.param .b32 _ZN8__main__2opB2v1B96cw51cXTLSUwv1sCUt9Ww0FEw09RRQPKiLTj0gIGIFp_2b2oLQFEYYkHSQB1OQAk0Bynm21OizQ1K0UoIGvDpQE8oxrNQE_3dEfff_param_3
+)
+{
+.reg .f32 	%f<7>;
+.reg .b32 	%r<2>;
+.reg .b64 	%rd<2>;
+
+
+ld.param.u64 	%rd1, [_ZN8__main__2opB2v1B96cw51cXTLSUwv1sCUt9Ww0FEw09RRQPKiLTj0gIGIFp_2b2oLQFEYYkHSQB1OQAk0Bynm21OizQ1K0UoIGvDpQE8oxrNQE_3dEfff_param_0];
+ld.param.f32 	%f1, [_ZN8__main__2opB2v1B96cw51cXTLSUwv1sCUt9Ww0FEw09RRQPKiLTj0gIGIFp_2b2oLQFEYYkHSQB1OQAk0Bynm21OizQ1K0UoIGvDpQE8oxrNQE_3dEfff_param_1];
+ld.param.f32 	%f2, [_ZN8__main__2opB2v1B96cw51cXTLSUwv1sCUt9Ww0FEw09RRQPKiLTj0gIGIFp_2b2oLQFEYYkHSQB1OQAk0Bynm21OizQ1K0UoIGvDpQE8oxrNQE_3dEfff_param_2];
+ld.param.f32 	%f3, [_ZN8__main__2opB2v1B96cw51cXTLSUwv1sCUt9Ww0FEw09RRQPKiLTj0gIGIFp_2b2oLQFEYYkHSQB1OQAk0Bynm21OizQ1K0UoIGvDpQE8oxrNQE_3dEfff_param_3];
+mul.f32 	%f4, %f1, %f3;
+sub.f32 	%f5, %f1, %f4;
+fma.rn.f32 	%f6, %f2, %f3, %f5;
+st.f32 	[%rd1], %f6;
+mov.u32 	%r1, 0;
+st.param.b32 	[func_retval0+0], %r1;
+ret;
+}
+)***";
+
+  float const LOW         = 100.0F;
+  float const HIGH        = 200.0F;
+  float const T           = 0.25F;
+  float const RES         = LOW - T * LOW + T * HIGH;
+  cudf::size_type const N = 200;
+
+  std::vector<float> low_host      = std::vector<float>(N, LOW);
+  std::vector<float> high_host     = std::vector<float>(N, HIGH);
+  std::vector<float> t_host        = std::vector<float>(N, T);
+  std::vector<float> t_scalar_host = std::vector<float>(1, T);
+  std::vector<float> expected_host = std::vector<float>(N, RES);
+
+  auto fourth()
+  {
+    return cudf::detail::make_counting_transform_iterator(0, [](auto i) { return (i % 4) == 0; });
+  }
+
+  auto fifth()
+  {
+    return cudf::detail::make_counting_transform_iterator(0, [](auto i) { return (i % 5) == 0; });
+  }
+};
+
+TEST_F(NullTest, ColumnNulls)
+{
+  auto low =
+    cudf::test::fixed_width_column_wrapper<float>(low_host.begin(), low_host.end(), fourth())
+      .release();
+  auto high =
+    cudf::test::fixed_width_column_wrapper<float>(high_host.begin(), high_host.end(), fifth())
+      .release();
+  auto t = cudf::test::fixed_width_column_wrapper<float>(t_host.begin(), t_host.end()).release();
+  auto expected =
+    cudf::test::fixed_width_column_wrapper<float>(expected_host.begin(), expected_host.end())
+      .release();
+
+  auto [expected_mask, expect_null_count] = cudf::bitmask_and(cudf::table_view({*low, *high}));
+
+  expected->set_null_mask(std::move(expected_mask), expect_null_count);
path: cpp/tests/transform/integration/unary_transform_test.cpp
user: login: ttnghia
body: Ignore if `null_count == 0`?
created_at: 2025-05-29T19:47:13Z
author_association: CONTRIBUTOR
diff_hunk: @@ -89,10 +89,10 @@ public class ColumnVectorTest extends CudfTestBase {
 
   @Test
   void testTransformVector() {
-    try (ColumnVector cv = ColumnVector.fromBoxedInts(2,3,4);
+    try (ColumnVector cv = ColumnVector.fromBoxedInts(2,3,null,4);
path: java/src/test/java/ai/rapids/cudf/ColumnVectorTest.java
user: login: ttnghia
body: ```suggestion
    try (ColumnVector cv = ColumnVector.fromBoxedInts(2, 3, null, 4);
```
created_at: 2025-05-29T19:48:09Z
author_association: CONTRIBUTOR
diff_hunk: @@ -89,10 +89,10 @@ public class ColumnVectorTest extends CudfTestBase {
 
   @Test
   void testTransformVector() {
-    try (ColumnVector cv = ColumnVector.fromBoxedInts(2,3,4);
+    try (ColumnVector cv = ColumnVector.fromBoxedInts(2,3,null,4);
          ColumnVector cv1 = cv.transform(ptx, true);
          ColumnVector cv2 = cv.transform(cuda, false);
-         ColumnVector expected = ColumnVector.fromBoxedInts(2*2-2, 3*3-3, 4*4-4)) {
+         ColumnVector expected = ColumnVector.fromBoxedInts(2*2-2, 3*3-3, null, 4*4-4)) {
path: java/src/test/java/ai/rapids/cudf/ColumnVectorTest.java
user: login: ttnghia
body: ```suggestion
         ColumnVector expected = ColumnVector.fromBoxedInts(2 * 2 - 2, 3 * 3 - 3, null, 4 * 4 - 4)) {
```
created_at: 2025-05-29T19:48:29Z
author_association: CONTRIBUTOR
diff_hunk: @@ -274,25 +279,58 @@ void launch_column_output_kernel(jitify2::ConfiguredKernel& kernel,
 template <typename T>
 void launch_span_kernel(jitify2::ConfiguredKernel& kernel,
                         device_span<T> output,
+                        rmm::device_buffer& null_mask,
                         std::vector<column_view> const& input_cols,
                         std::optional<void*> user_data,
                         rmm::cuda_stream_view stream,
                         rmm::device_async_resource_ref mr)
 {
-  auto outputs = to_device_vector(
-    std::vector{cudf::jit::device_span<T>{output.data(), output.size()}}, stream, mr);
+  auto outputs = to_device_vector(std::vector{cudf::jit::device_optional_span<T>{
+                                    cudf::jit::device_span<T>{output.data(), output.size()},
+                                    static_cast<bitmask_type*>(null_mask.data())}},
+                                  stream,
+                                  mr);
+
   auto [input_handles, inputs] =
     column_views_to_device<column_device_view, column_view>(input_cols, stream, mr);
 
-  cudf::jit::device_span<T> const* outputs_ptr = outputs.data();
-  column_device_view const* inputs_ptr         = inputs.data();
-  void* p_user_data                            = user_data.value_or(nullptr);
+  cudf::jit::device_optional_span<T> const* outputs_ptr = outputs.data();
path: cpp/src/transform/transform.cpp
user: login: lamarrr
body: yes
created_at: 2025-05-29T19:54:16Z
author_association: CONTRIBUTOR
diff_hunk: @@ -302,14 +340,10 @@ std::unique_ptr<column> transform_operation(column_view base_column,
                                             rmm::cuda_stream_view stream,
                                             rmm::device_async_resource_ref mr)
 {
-  auto output = make_fixed_width_column(output_type,
-                                        base_column.size(),
-                                        copy_bitmask(base_column, stream, mr),
-                                        base_column.null_count(),
-                                        stream,
-                                        mr);
+  auto [null_mask, null_count] = make_transform_null_mask(base_column, inputs, stream, mr);
 
-  if (base_column.is_empty()) { return output; }
+  auto output = make_fixed_width_column(
+    output_type, base_column.size(), std::move(null_mask), null_count, stream, mr);
path: cpp/src/transform/transform.cpp
user: login: lamarrr
body: https://github.com/rapidsai/cudf/pull/18845/files#diff-5bf667c1173053135dd69543fe01a13aec1b5dece7f3de53a9d5a0fd5f3eb495R331

that function uses the `cudf::bitmask_and` function, it already sets the null-mask to nullptr/empty state if the null-count is 0.

https://github.com/rapidsai/cudf/blob/2c4f09fd25aef591fd71ef0aa315c93ba4c5bbb8/cpp/include/cudf/null_mask.hpp#L182
created_at: 2025-05-29T19:57:42Z
author_association: CONTRIBUTOR
diff_hunk: @@ -350,11 +384,11 @@ std::unique_ptr<column> string_view_operation(column_view base_column,
 
   rmm::device_uvector<string_view> string_views(base_column.size(), stream, mr);
 
-  launch_span_kernel<string_view>(kernel, string_views, inputs, user_data, stream, mr);
+  launch_span_kernel<string_view>(kernel, string_views, null_mask, inputs, user_data, stream, mr);
 
   auto output = make_strings_column(string_views, string_view{}, stream, mr);
 
-  output->set_null_mask(copy_bitmask(base_column, stream, mr), base_column.null_count());
+  output->set_null_mask(std::move(null_mask), null_count);
path: cpp/src/transform/transform.cpp
user: login: lamarrr
body: same as above
created_at: 2025-05-29T19:58:11Z
author_association: CONTRIBUTOR
diff_hunk: @@ -640,4 +648,197 @@ __device__ void transform(void* user_data, cudf::size_type row,
   CUDF_TEST_EXPECT_COLUMNS_EQUAL(expected, result->view());
 }
 
+struct NullTest : public cudf::test::BaseFixture {
+ protected:
+  void SetUp() override
+  {
+    if (!cudf::is_runtime_jit_supported()) {
+      GTEST_SKIP() << "Skipping tests that require runtime JIT support";
+    }
+  }
+
+  char const* const cuda =
+    "__device__ inline void lerp(float* output, float low, float high, float t){*output = low - t "
+    "* low + t * high; }";
+
+  // Generated from NUMBA, using:
+  //
+  // ```py
+  //
+  // from numba import cuda, float32
+  // from numba.cuda import compile_ptx_for_current_device
+  //
+  // # Define a CUDA device function
+  //
+  // @cuda.jit(device=True)
+  // def op(low, high, t):
+  //         return low - t * low + t * high
+  //
+  // # Define argument types for the function
+  // arg_types = (float32, float32, float32)
+  //
+  // # Compile the device function as relocatable
+  // ptx, _ = cuda.compile_ptx_for_current_device(op, arg_types, device=True)
+  //
+  //
+  // # Print the PTX code
+  // print("Relocatable PTX Code:")
+  // print(ptx)
+  //
+  //
+  // ```
+  //
+  char const* const ptx =
+    R"***(
+//
+// Generated by NVIDIA NVVM Compiler
+//
+// Compiler Build ID: CL-35813241
+// Cuda compilation tools, release 12.9, V12.9.41
+// Based on NVVM 7.0.1
+//
+
+.version 8.8
+.target sm_86
+.address_size 64
+
+// .globl	_ZN8__main__2opB2v1B96cw51cXTLSUwv1sCUt9Ww0FEw09RRQPKiLTj0gIGIFp_2b2oLQFEYYkHSQB1OQAk0Bynm21OizQ1K0UoIGvDpQE8oxrNQE_3dEfff
+.common .global .align 8 .u64 _ZN08NumbaEnv8__main__2opB2v1B96cw51cXTLSUwv1sCUt9Ww0FEw09RRQPKiLTj0gIGIFp_2b2oLQFEYYkHSQB1OQAk0Bynm21OizQ1K0UoIGvDpQE8oxrNQE_3dEfff;
+
+.visible .func  (.param .b32 func_retval0) _ZN8__main__2opB2v1B96cw51cXTLSUwv1sCUt9Ww0FEw09RRQPKiLTj0gIGIFp_2b2oLQFEYYkHSQB1OQAk0Bynm21OizQ1K0UoIGvDpQE8oxrNQE_3dEfff(
+.param .b64 _ZN8__main__2opB2v1B96cw51cXTLSUwv1sCUt9Ww0FEw09RRQPKiLTj0gIGIFp_2b2oLQFEYYkHSQB1OQAk0Bynm21OizQ1K0UoIGvDpQE8oxrNQE_3dEfff_param_0,
+.param .b32 _ZN8__main__2opB2v1B96cw51cXTLSUwv1sCUt9Ww0FEw09RRQPKiLTj0gIGIFp_2b2oLQFEYYkHSQB1OQAk0Bynm21OizQ1K0UoIGvDpQE8oxrNQE_3dEfff_param_1,
+.param .b32 _ZN8__main__2opB2v1B96cw51cXTLSUwv1sCUt9Ww0FEw09RRQPKiLTj0gIGIFp_2b2oLQFEYYkHSQB1OQAk0Bynm21OizQ1K0UoIGvDpQE8oxrNQE_3dEfff_param_2,
+.param .b32 _ZN8__main__2opB2v1B96cw51cXTLSUwv1sCUt9Ww0FEw09RRQPKiLTj0gIGIFp_2b2oLQFEYYkHSQB1OQAk0Bynm21OizQ1K0UoIGvDpQE8oxrNQE_3dEfff_param_3
+)
+{
+.reg .f32 	%f<7>;
+.reg .b32 	%r<2>;
+.reg .b64 	%rd<2>;
+
+
+ld.param.u64 	%rd1, [_ZN8__main__2opB2v1B96cw51cXTLSUwv1sCUt9Ww0FEw09RRQPKiLTj0gIGIFp_2b2oLQFEYYkHSQB1OQAk0Bynm21OizQ1K0UoIGvDpQE8oxrNQE_3dEfff_param_0];
+ld.param.f32 	%f1, [_ZN8__main__2opB2v1B96cw51cXTLSUwv1sCUt9Ww0FEw09RRQPKiLTj0gIGIFp_2b2oLQFEYYkHSQB1OQAk0Bynm21OizQ1K0UoIGvDpQE8oxrNQE_3dEfff_param_1];
+ld.param.f32 	%f2, [_ZN8__main__2opB2v1B96cw51cXTLSUwv1sCUt9Ww0FEw09RRQPKiLTj0gIGIFp_2b2oLQFEYYkHSQB1OQAk0Bynm21OizQ1K0UoIGvDpQE8oxrNQE_3dEfff_param_2];
+ld.param.f32 	%f3, [_ZN8__main__2opB2v1B96cw51cXTLSUwv1sCUt9Ww0FEw09RRQPKiLTj0gIGIFp_2b2oLQFEYYkHSQB1OQAk0Bynm21OizQ1K0UoIGvDpQE8oxrNQE_3dEfff_param_3];
+mul.f32 	%f4, %f1, %f3;
+sub.f32 	%f5, %f1, %f4;
+fma.rn.f32 	%f6, %f2, %f3, %f5;
+st.f32 	[%rd1], %f6;
+mov.u32 	%r1, 0;
+st.param.b32 	[func_retval0+0], %r1;
+ret;
+}
+)***";
+
+  float const LOW         = 100.0F;
+  float const HIGH        = 200.0F;
+  float const T           = 0.25F;
+  float const RES         = LOW - T * LOW + T * HIGH;
+  cudf::size_type const N = 200;
+
+  std::vector<float> low_host      = std::vector<float>(N, LOW);
+  std::vector<float> high_host     = std::vector<float>(N, HIGH);
+  std::vector<float> t_host        = std::vector<float>(N, T);
+  std::vector<float> t_scalar_host = std::vector<float>(1, T);
+  std::vector<float> expected_host = std::vector<float>(N, RES);
+
+  auto fourth()
+  {
+    return cudf::detail::make_counting_transform_iterator(0, [](auto i) { return (i % 4) == 0; });
+  }
+
+  auto fifth()
+  {
+    return cudf::detail::make_counting_transform_iterator(0, [](auto i) { return (i % 5) == 0; });
+  }
+};
+
+TEST_F(NullTest, ColumnNulls)
+{
+  auto low =
+    cudf::test::fixed_width_column_wrapper<float>(low_host.begin(), low_host.end(), fourth())
+      .release();
+  auto high =
+    cudf::test::fixed_width_column_wrapper<float>(high_host.begin(), high_host.end(), fifth())
+      .release();
+  auto t = cudf::test::fixed_width_column_wrapper<float>(t_host.begin(), t_host.end()).release();
+  auto expected =
+    cudf::test::fixed_width_column_wrapper<float>(expected_host.begin(), expected_host.end())
+      .release();
+
+  auto [expected_mask, expect_null_count] = cudf::bitmask_and(cudf::table_view({*low, *high}));
+
+  expected->set_null_mask(std::move(expected_mask), expect_null_count);
path: cpp/tests/transform/integration/unary_transform_test.cpp
user: login: lamarrr
body: same as above
created_at: 2025-05-29T19:59:37Z
author_association: CONTRIBUTOR
diff_hunk: @@ -274,25 +279,58 @@ void launch_column_output_kernel(jitify2::ConfiguredKernel& kernel,
 template <typename T>
 void launch_span_kernel(jitify2::ConfiguredKernel& kernel,
                         device_span<T> output,
+                        rmm::device_buffer& null_mask,
                         std::vector<column_view> const& input_cols,
                         std::optional<void*> user_data,
                         rmm::cuda_stream_view stream,
                         rmm::device_async_resource_ref mr)
 {
-  auto outputs = to_device_vector(
-    std::vector{cudf::jit::device_span<T>{output.data(), output.size()}}, stream, mr);
+  auto outputs = to_device_vector(std::vector{cudf::jit::device_optional_span<T>{
+                                    cudf::jit::device_span<T>{output.data(), output.size()},
+                                    static_cast<bitmask_type*>(null_mask.data())}},
+                                  stream,
+                                  mr);
+
   auto [input_handles, inputs] =
     column_views_to_device<column_device_view, column_view>(input_cols, stream, mr);
 
-  cudf::jit::device_span<T> const* outputs_ptr = outputs.data();
-  column_device_view const* inputs_ptr         = inputs.data();
-  void* p_user_data                            = user_data.value_or(nullptr);
+  cudf::jit::device_optional_span<T> const* outputs_ptr = outputs.data();
path: cpp/src/transform/transform.cpp
user: login: lamarrr
body: for context: the JIT ABI can't pass structs directly (especially if it is larger than a pointer), it needs pointers (just like any ABI/FFI)
created_at: 2025-05-29T20:07:53Z
author_association: CONTRIBUTOR

PR comments:
user: login: copy-pr-bot[bot]
created_at: 2025-05-15T18:24:56Z
body: Auto-sync is disabled for draft pull requests in this repository. Workflows must be run manually.

Contributors can view more details about this message [here](https://docs.gha-runners.nvidia.com/cpr/auto-sync).

author_association: NONE
user: login: lamarrr
created_at: 2025-05-19T13:03:19Z
body: /ok to test d8b0433
author_association: CONTRIBUTOR
user: login: lamarrr
created_at: 2025-05-19T13:08:41Z
body: /ok to test 8dded0a
author_association: CONTRIBUTOR
user: login: lamarrr
created_at: 2025-05-19T20:25:15Z
body: /ok to test 08ecef2
author_association: CONTRIBUTOR
user: login: lamarrr
created_at: 2025-05-29T21:48:36Z
body: /merge
author_association: CONTRIBUTOR