[FEA] Improve Null-Aware Operator Support in AST-Codegen

cpp/CMakeLists.txt

@@ -824,7 +824,7 @@ add_library(
   src/transform/nans_to_nulls.cu
   src/transform/one_hot_encode.cu
   src/transform/row_bit_count.cu
-  src/transform/transform.cpp
+  src/transform/transform.cu
   src/transpose/transpose.cu
   src/unary/cast_ops.cu
   src/unary/math_ops.cu
@@ -854,7 +854,7 @@ set_source_files_properties(
   src/rolling/detail/rolling_variable_window.cu
   src/rolling/grouped_rolling.cu
   src/rolling/rolling.cu
-  src/transform/transform.cpp
+  src/transform/transform.cu
   PROPERTIES COMPILE_DEFINITIONS "_FILE_OFFSET_BITS=64"
 )
 

cpp/benchmarks/ndsh/q09.cpp

@@ -164,6 +164,7 @@ struct q9_data {
                          false,
                          std::nullopt,
                          cudf::null_aware::NO,
+                         cudf::null_output::PRESERVE,
                          stream,
                          mr);
 }

cpp/benchmarks/transform/polynomials.cpp

@@ -96,6 +96,7 @@ static void BM_transform_polynomials(nvbench::state& state)
                     false,
                     std::nullopt,
                     cudf::null_aware::NO,
+                    cudf::null_output::PRESERVE,
                     launch.get_stream().get_stream());
   });
 }

cpp/benchmarks/transform/transform.cpp

@@ -102,6 +102,7 @@ static void BM_transform(nvbench::state& state)
                     false,
                     std::nullopt,
                     cudf::null_aware::NO,
+                    cudf::null_output::PRESERVE,
                     launch.get_stream().get_stream());
   });
 }

cpp/examples/string_transforms/compute_checksum_jit.cpp

@@ -55,6 +55,7 @@ std::tuple<std::unique_ptr<cudf::column>, std::vector<int32_t>> transform(
                                 false,
                                 std::nullopt,
                                 cudf::null_aware::NO,
+                                cudf::null_output::PRESERVE,
                                 stream,
                                 mr);
 

cpp/examples/string_transforms/extract_email_jit.cpp

@@ -70,6 +70,7 @@ __device__ void email_provider(cudf::string_view* out,
                                    false,
                                    std::nullopt,
                                    cudf::null_aware::NO,
+                                   cudf::null_output::PRESERVE,
                                    stream,
                                    mr);
 

cpp/examples/string_transforms/format_phone_jit.cpp

@@ -133,6 +133,7 @@ __device__ void e164_format(void* scratch,
                     false,
                     scratch.data(),
                     cudf::null_aware::NO,
+                    cudf::null_output::PRESERVE,
                     stream,
                     mr);
 

cpp/examples/string_transforms/localize_phone_jit.cpp

@@ -156,6 +156,7 @@ __device__ void format_phone(void* scratch,
                                 false,
                                 scratch.data(),
                                 cudf::null_aware::NO,
+                                cudf::null_output::PRESERVE,
                                 stream,
                                 mr);
 

cpp/include/cudf/transform.hpp

@@ -57,6 +57,7 @@ namespace CUDF_EXPORT cudf {
  * @param is_ptx        true: the UDF is treated as PTX code; false: the UDF is treated as CUDA code
  * @param user_data     User-defined device data to pass to the UDF.
  * @param is_null_aware Signifies the UDF will receive row inputs as optional values
+ * @param null_policy   Signifies if a null mask should be created for the output column
  * @param stream        CUDA stream used for device memory operations and kernel launches
  * @param mr            Device memory resource used to allocate the returned column's device memory
  * @return              The column resulting from applying the transform function to
@@ -69,6 +70,7 @@ std::unique_ptr<column> transform(
   bool is_ptx,
   std::optional<void*> user_data    = std::nullopt,
   null_aware is_null_aware          = null_aware::NO,
+  null_output null_policy           = null_output::PRESERVE,
   rmm::cuda_stream_view stream      = cudf::get_default_stream(),
   rmm::device_async_resource_ref mr = cudf::get_current_device_resource_ref());
 

cpp/include/cudf/types.hpp

@@ -240,6 +240,12 @@ enum class null_aware : bool {
   YES = 1   ///< The function is null-aware
 };
 
+/// @brief Indicates the null output policy of a function.
+enum class null_output : uint8_t {
+  PRESERVE     = 0,  ///< The resulting column's nullmask is preserved
+  NON_NULLABLE = 1   ///< A null-mask is not produced and all values are considered valid
+};
+
 /**
  * @brief Indicator for the logical data type of an element in a column.
  *

cpp/src/jit/accessors.cuh

@@ -52,6 +52,12 @@ struct column_accessor {
     return inputs[index].is_null(row);
   }
 
+  template <typename ColumnView>
+  static __device__ bool is_valid(ColumnView const* inputs, cudf::size_type row)
+  {
+    return inputs[index].is_valid(row);
+  }
+
   template <typename ColumnView>
   static __device__ cuda::std::optional<T> nullable_element(ColumnView const* columns,
                                                             cudf::size_type row)
@@ -85,6 +91,12 @@ struct span_accessor {
     return inputs[index].is_null(row);
   }
 
+  static __device__ bool is_valid(cudf::jit::device_optional_span<T> const* inputs,
+                                  cudf::size_type row)
+  {
+    return inputs[index].is_valid(row);
+  }
+
   static __device__ cuda::std::optional<T> nullable_element(
     cudf::jit::device_optional_span<T> const* outputs, cudf::size_type row)
   {
@@ -117,6 +129,12 @@ struct scalar_accessor {
     return Accessor::is_null(columns, 0);
   }
 
+  template <typename ColumnView>
+  static __device__ bool is_valid(ColumnView const* columns, cudf::size_type)
+  {
+    return Accessor::is_valid(columns, 0);
+  }
+
   template <typename ColumnView>
   static __device__ decltype(auto) nullable_element(ColumnView const* columns, cudf::size_type)
   {

cpp/src/jit/helpers.cpp

@@ -39,42 +39,6 @@ typename std::vector<column_view>::const_iterator get_transform_base_column(
   return largest;
 }
 
-jitify2::StringVec build_jit_template_params(
-  bool has_user_data,
-  null_aware is_null_aware,
-  std::vector<std::string> const& span_outputs,
-  std::vector<std::string> const& column_outputs,
-  std::vector<input_column_reflection> const& column_inputs)
-{
-  jitify2::StringVec tparams;
-
-  tparams.emplace_back(jitify2::reflection::reflect(has_user_data));
-  tparams.emplace_back(jitify2::reflection::reflect(is_null_aware == null_aware::YES));
-
-  std::transform(thrust::counting_iterator<size_t>(0),
-                 thrust::counting_iterator(span_outputs.size()),
-                 std::back_inserter(tparams),
-                 [&](auto i) {
-                   return jitify2::reflection::Template("cudf::jit::span_accessor")
-                     .instantiate(span_outputs[i], i);
-                 });
-
-  std::transform(thrust::counting_iterator<size_t>(0),
-                 thrust::counting_iterator(column_outputs.size()),
-                 std::back_inserter(tparams),
-                 [&](auto i) {
-                   return jitify2::reflection::Template("cudf::jit::column_accessor")
-                     .instantiate(column_outputs[i], i);
-                 });
-
-  std::transform(thrust::counting_iterator<size_t>(0),
-                 thrust::counting_iterator(column_inputs.size()),
-                 std::back_inserter(tparams),
-                 [&](auto i) { return column_inputs[i].accessor(i); });
-
-  return tparams;
-}
-
 std::map<uint32_t, std::string> build_ptx_params(std::vector<std::string> const& output_typenames,
                                                  std::vector<std::string> const& input_typenames,
                                                  bool has_user_data)

cpp/src/jit/helpers.hpp

@@ -34,6 +34,7 @@ constexpr bool is_scalar(cudf::size_type base_column_size, cudf::size_type colum
 
 typename std::vector<column_view>::const_iterator get_transform_base_column(
   std::vector<column_view> const& inputs);
+
 struct input_column_reflection {
   std::string type_name;
   bool is_scalar = false;
@@ -49,13 +50,6 @@ struct input_column_reflection {
   }
 };
 
-jitify2::StringVec build_jit_template_params(
-  bool has_user_data,
-  null_aware is_null_aware,
-  std::vector<std::string> const& span_outputs,
-  std::vector<std::string> const& column_outputs,
-  std::vector<input_column_reflection> const& column_inputs);
-
 std::map<uint32_t, std::string> build_ptx_params(std::vector<std::string> const& output_typenames,
                                                  std::vector<std::string> const& input_typenames,
                                                  bool has_user_data);