[FEA] Support input of dictionary-encoded types in Transforms

cpp/include/cudf/column/column_child_offsets.hpp

@@ -0,0 +1,25 @@
+/*
+ * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * SPDX-License-Identifier: Apache-2.0
+ */
+#pragma once
+
+/**
+ * @file column_child_offsets.hpp
+ * @brief Constants for child column indices within compound column types
+ */
+
+#include <cudf/types.hpp>
+#include <cudf/utilities/export.hpp>
+
+namespace CUDF_EXPORT cudf {
+
+static constexpr size_type offsets_column_index = 0;  ///< Child index of the offsets column
+
+static constexpr size_type dictionary_indices_column_index =
+  0;  ///< Child index of the dictionary indices column
+
+static constexpr size_type dictionary_keys_column_index =
+  1;  ///< Child index of the dictionary key column
+
+}  // namespace CUDF_EXPORT cudf

cpp/include/cudf/column/column_device_view_base.cuh

@@ -1,9 +1,10 @@
 /*
- * SPDX-FileCopyrightText: Copyright (c) 2019-2026, NVIDIA CORPORATION.
+ * SPDX-FileCopyrightText: Copyright (c) 2019-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
  * SPDX-License-Identifier: Apache-2.0
  */
 #pragma once
 
+#include <cudf/column/column_child_offsets.hpp>
 #include <cudf/detail/offsets_iterator.cuh>
 #include <cudf/fixed_point/fixed_point.hpp>
 #include <cudf/strings/string_view.cuh>
@@ -61,15 +62,45 @@ struct nullate {
   };
 };
 
+/**
+ * @brief A type tag to specify that a column should be treated as a dictionary column
+ *
+ * @tparam IndexType The type of the dictionary indices
+ * @tparam KeyType The type of the dictionary keys
+ */
+template <typename IndexType, typename KeyType>
+  requires(is_index_type<IndexType>() && is_relationally_comparable<KeyType, KeyType>())
+struct dictionary_element {
+  using index_type = IndexType;  ///< The type of the dictionary indices
+  using key_type   = KeyType;    ///< The type of the dictionary keys
+
+  key_type key{};  ///< The dictionary key for this element
+};
+
+/**
+ * @brief A type trait to determine if a type is a dictionary encoded type.
+ * @tparam T The type to check
+ */
+template <typename T>
+inline constexpr bool is_dictionary_encoded = false;
+
+/**
+ * @brief A type trait to determine if a type is a dictionary encoded type.
+ * @tparam IndexType The type of the dictionary indices
+ * @tparam KeyType The type of the dictionary keys
+ */
+template <typename IndexType, typename KeyType>
+inline constexpr bool is_dictionary_encoded<dictionary_element<IndexType, KeyType>> = true;
+
 namespace detail {
 /**
  * @brief An immutable, non-owning view of device data as a column of elements
  * that is trivially copyable and usable in CUDA device code.
  */
 class alignas(16) column_device_view_base {
  public:
-  // TODO: merge this offsets column index with `strings_column_view::offsets_column_index`
-  static constexpr size_type offsets_column_index{0};  ///< Child index of the offsets column
+  static constexpr size_type offsets_column_index =
+    cudf::offsets_column_index;  ///< Child index of the offsets column
 
   column_device_view_base()                               = delete;
   ~column_device_view_base()                              = default;
@@ -464,6 +495,33 @@ class alignas(16) column_device_view_core : public detail::column_device_view_ba
     return T{scaled_integer<rep>{data<rep>()[element_index], scale}};
   }
 
+  /**
+   * @brief Returns a decoded copy of the element at the specified index.
+   *
+   * If the element at the specified index is NULL, i.e.,
+   * `is_null(element_index) == true`, then any attempt to use the result will
+   * lead to undefined behavior.
+   *
+   * This function accounts for the offset.
+   *
+   * This function does not participate in overload resolution if `is_dictionary_encoded<T>` is
+   * false.
+   *
+   * @tparam T The element type, e.g., `dictionary_element<int32_t, float>` for a dictionary column
+   * with `int32_t` indices and `float` keys
+   * @param element_index Position of the desired element
+   * @return The element at the specified index
+   */
+  template <typename T, CUDF_ENABLE_IF(is_dictionary_encoded<T>)>
+  [[nodiscard]] __device__ decltype(auto) element(size_type element_index) const noexcept
+  {
+    auto const& indices = child(dictionary_indices_column_index);
+    auto const& keys    = child(dictionary_keys_column_index);
+    auto const index    = indices.template element<typename T::index_type>(
+      element_index + offset());  // account for this view's _offset
+    return keys.template element<typename T::key_type>(index);
+  }
+
   /**
    * @brief Returns a nullable element at the specified index. If the element is null, returns
    * `nullopt`.
@@ -538,7 +596,8 @@ class alignas(16) column_device_view_core : public detail::column_device_view_ba
 
 /**
  * @brief A mutable, non-owning view of device data as a column of elements
- * that is trivially copyable and usable in CUDA device code and offline-compiled code (i.e. NVRTC).
+ * that is trivially copyable and usable in CUDA device code and offline-compiled code (i.e.
+ * NVRTC).
  *
  * @ingroup column_classes
  */

cpp/include/cudf/dictionary/dictionary_column_view.hpp

@@ -1,9 +1,10 @@
 /*
- * SPDX-FileCopyrightText: Copyright (c) 2020-2024, NVIDIA CORPORATION.
+ * SPDX-FileCopyrightText: Copyright (c) 2020-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
  * SPDX-License-Identifier: Apache-2.0
  */
 #pragma once
 
+#include <cudf/column/column_child_offsets.hpp>
 #include <cudf/column/column_view.hpp>
 
 /**
@@ -52,9 +53,9 @@ class dictionary_column_view : private column_view {
   dictionary_column_view& operator=(dictionary_column_view&&) = default;
 
   /// Index of the indices column of the dictionary column
-  static constexpr size_type indices_column_index{0};
+  static constexpr size_type indices_column_index = cudf::dictionary_indices_column_index;
   /// Index of the keys column of the dictionary column
-  static constexpr size_type keys_column_index{1};
+  static constexpr size_type keys_column_index = cudf::dictionary_keys_column_index;
 
   using column_view::has_nulls;
   using column_view::is_empty;

cpp/include/cudf/transform.hpp

@@ -95,9 +95,11 @@ struct transform_output {
  *
  *
  * @throws std::invalid_argument if any of the input columns have different sizes (except scalars)
- * @throws std::invalid_argument if `output_type` or any of the input types are not supported.
- * CUDA-supported types are fixed-width and string types, while PTX-supported types are integral and
- * floating-point types.
+ * @throws std::invalid_argument if any of the output or input types are not supported.
+ * CUDA-supported input types are fixed-width, string, and their dictionary types. PTX-supported
+ * input types are integrals, floats, and their dictionary types. CUDA-supported output types are
+ * fixed-width, string, and their dictionary types. PTX-supported output types are integrals,
+ * floats, and their dictionary types.
  * @throws std::invalid_argument if the inputs only have a scalar with no column inputs and
  * `row_size` is not provided. This is because the row size cannot be inferred from the inputs in
  * this case.
@@ -145,6 +147,11 @@ std::unique_ptr<column> transform_extended(
  * @throws std::invalid_argument if the inputs only have a scalar with no column inputs and
  * `row_size` is not provided. This is because the row size cannot be inferred from the inputs in
  * this case.
+ * @throws std::invalid_argument if any of the output or input types are not supported.
+ * CUDA-supported input types are fixed-width, string, and their dictionary types. PTX-supported
+ * input types are integrals, floats, and their dictionary types. CUDA-supported output types are
+ * fixed-width, string, and their dictionary types. PTX-supported output types are integrals,
+ * floats, and their dictionary types.
  * @throws std::invalid_argument if string offsets are provided for non-string output columns, or
  * if the number of string offsets does not match the number of output columns.
  *

cpp/src/jit/column_accessor.cuh

@@ -1,5 +1,5 @@
 /*
- * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION.
+ * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
  * SPDX-License-Identifier: Apache-2.0
  */
 
@@ -45,7 +45,7 @@ struct column_accessor {
     return reinterpret_cast<column_type const&>(cols[index]);
   }
 
-  static __device__ element_type element(auto const* __restrict__ cols, size_type row)
+  static __device__ auto element(auto const* __restrict__ cols, size_type row)
   {
     return column(cols).template element<element_type>(map_index(row));
   }
@@ -60,8 +60,7 @@ struct column_accessor {
     return column(cols).is_valid(map_index(row));
   }
 
-  static __device__ optional_element_type nullable_element(auto const* __restrict__ cols,
-                                                           size_type row)
+  static __device__ auto nullable_element(auto const* __restrict__ cols, size_type row)
   {
     return column(cols).template nullable_element<element_type>(map_index(row));
   }

cpp/src/jit/helpers.cpp

@@ -1,16 +1,20 @@
 /*
- * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION.
+ * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
  * SPDX-License-Identifier: Apache-2.0
  */
 
 #include "helpers.hpp"
 
+#include <cudf/column/column_device_view_base.cuh>
 #include <cudf/detail/nvtx/ranges.hpp>
+#include <cudf/utilities/type_dispatcher.hpp>
 
 #include <jit/cache.hpp>
 #include <rtcx.hpp>
 #include <runtime/context.hpp>
 
+#include <format>
+
 namespace cudf {
 namespace jit {
 
@@ -77,18 +81,6 @@ std::map<uint32_t, std::string> build_ptx_params(std::span<std::string const> ou
   return params;
 }
 
-std::vector<std::string> input_type_names(
-  std::span<std::variant<column_view, scalar_column_view> const> views)
-{
-  std::vector<std::string> names;
-
-  std::transform(views.begin(), views.end(), std::back_inserter(names), [&](auto const& view) {
-    return std::visit([](auto& a) { return type_to_name(a.type()); }, view);
-  });
-
-  return names;
-}
-
 kernel get_udf_kernel(std::string const& source_file,
                       std::string const& kernel_name,
                       std::string const& cuda_source)

cpp/src/jit/helpers.hpp

@@ -1,5 +1,5 @@
 /*
- * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION.
+ * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
  * SPDX-License-Identifier: Apache-2.0
  */
 #pragma once
@@ -72,9 +72,6 @@ column_views_to_device(std::span<ColumnView const> views,
   return std::make_tuple(std::move(handles), std::move(device_array));
 }
 
-std::vector<std::string> input_type_names(
-  std::span<std::variant<column_view, scalar_column_view> const> views);
-
 kernel get_udf_kernel(std::string const& source_file,
                       std::string const& kernel_name,
                       std::string const& cuda_source);