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

cpp/include/cudf/column/column_child_offsets.h

@@ -0,0 +1,20 @@
+/*
+ * SPDX-FileCopyrightText: Copyright (c) 2019-2026, NVIDIA CORPORATION.
+ * SPDX-License-Identifier: Apache-2.0
+ */
+#pragma once
+
+#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 offsets 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

@@ -4,6 +4,7 @@
  */
 #pragma once
 
+#include <cudf/column/column_child_offsets.h>
 #include <cudf/detail/offsets_iterator.cuh>
 #include <cudf/fixed_point/fixed_point.hpp>
 #include <cudf/strings/string_view.cuh>
@@ -63,15 +64,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;
@@ -447,6 +478,32 @@ 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);
+    return keys.template element<typename T::key_type>(index);
+  }
+
   /**
    * @brief Returns the specified child
    *

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.
  * SPDX-License-Identifier: Apache-2.0
  */
 #pragma once
 
+#include <cudf/column/column_child_offsets.h>
 #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/src/jit/helpers.cpp

@@ -5,10 +5,14 @@
 
 #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 <format>
+
 namespace cudf {
 namespace jit {
 
@@ -111,13 +115,58 @@ std::vector<std::string> input_type_names(
   return names;
 }
 
-input_reflection reflect_input(std::variant<column_view, scalar_column_view> const& input)
+namespace {
+
+std::string get_jit_element_type_name_impl(column_view const& view);
+
+struct jit_element_type_name_fn {
+  template <typename T>
+    requires(is_fixed_width<T>() || std::is_same_v<T, cudf::string_view>)
+  std::string operator()(column_view const& view) const
+  {
+    return type_to_name(view.type());
+  }
+
+  template <typename T>
+    requires(std::is_same_v<T, cudf::dictionary32>)
+  std::string operator()(column_view const& view) const
+  {
+    return std::format(
+      "cudf::dictionary_element<{}, {}>",
+      get_jit_element_type_name_impl(view.child(cudf::dictionary_indices_column_index)),
+      get_jit_element_type_name_impl(view.child(cudf::dictionary_keys_column_index)));
+  }
+
+  template <typename T>
+    requires(!is_fixed_width<T>() && !std::is_same_v<T, cudf::string_view> &&
+             !std::is_same_v<T, cudf::dictionary32>)
+  std::string operator()(column_view const& view) const
+  {
+    CUDF_FAIL("Unsupported type for JIT compilation: " + type_to_name(view.type()));
+  }
+};
+
+std::string get_jit_element_type_name_impl(column_view const& view)
 {
-  auto get_type_name = [](auto const& var) {
-    return std::visit([](auto& a) { return type_to_name(a.type()); }, var);
-  };
+  return cudf::type_dispatcher(view.type(), jit_element_type_name_fn{}, view);
+}
+
+std::string get_jit_element_type_name(column_view const& view)
+{
+  return get_jit_element_type_name_impl(view);
+}
 
-  return input_reflection{get_type_name(input), std::holds_alternative<scalar_column_view>(input)};
+std::string get_jit_element_type_name(scalar_column_view const& view)
+{
+  return get_jit_element_type_name(view.as_column_view());
+}
+
+}  // namespace
+
+input_reflection reflect_input(std::variant<column_view, scalar_column_view> const& input)
+{
+  auto type_name = std::visit([](auto& a) { return get_jit_element_type_name(a); }, input);
+  return input_reflection{type_name, std::holds_alternative<scalar_column_view>(input)};
 }
 
 std::vector<input_reflection> reflect_inputs(

cpp/src/transform/transform.cu

@@ -427,12 +427,14 @@ void perform_checks(std::optional<size_type> in_row_size, data_type output_type,
     return std::visit([](auto const& col) { return col.type(); }, in);
   };
 
-  CUDF_EXPECTS(
-    std::all_of(thrust::make_transform_iterator(inputs.begin(), get_type),
-                thrust::make_transform_iterator(inputs.end(), get_type),
-                [](data_type t) { return is_fixed_width(t) || (t.id() == type_id::STRING); }),
-    "Transforms only support input of fixed-width or string types",
-    std::invalid_argument);
+  CUDF_EXPECTS(std::all_of(thrust::make_transform_iterator(inputs.begin(), get_type),
+                           thrust::make_transform_iterator(inputs.end(), get_type),
+                           [](data_type t) {
+                             return is_fixed_width(t) || (t.id() == type_id::STRING) ||
+                                    cudf::is_dictionary(t);
+                           }),
+               "Transforms only support input of fixed-width, string, and dictionary types",
+               std::invalid_argument);
 
   check_row_size(in_row_size, inputs);
 }

cpp/tests/transform/integration/unary_transform_test.cpp

@@ -29,6 +29,7 @@
 #include <cudf_test/type_lists.hpp>
 
 #include <cudf/detail/iterator.cuh>
+#include <cudf/dictionary/encode.hpp>
 #include <cudf/transform.hpp>
 
 namespace transformation {
@@ -313,6 +314,89 @@ __device__ inline void f(cudf::timestamp_us* output, cudf::timestamp_us input)
   test_udf<dtype>(cuda, op, data_init, 500, cudf::udf_source_type::CUDA);
 }
 
+TEST_F(UnaryOperationIntegrationTest, Transform_DictionaryString)
+{
+  std::string const cuda =
+    R"***(
+__device__ inline void decode(cudf::string_view * output, cudf::string_view input){
+  *output = input;
+})***";
+
+  // non-nullable
+  {
+    auto a =
+      cudf::test::strings_column_wrapper{
+        "eee", "aaa", "ddd", "bbb", "ccc", "ccc", "ccc", "eee", "aaa"}
+        .release();
+
+    auto a_encoded = cudf::dictionary::encode(a->view());
+
+    cudf::transform_input inputs[] = {*a_encoded};
+
+    auto out =
+      cudf::transform_extended(inputs, cuda, cudf::data_type{cudf::type_id::STRING}, false);
+
+    CUDF_TEST_EXPECT_COLUMNS_EQUAL(out->view(), a->view());
+  }
+
+  // nullable
+  {
+    auto a =
+      cudf::test::strings_column_wrapper{
+        {"eee", "aaa", "ddd", "bbb", "ccc", "ccc", "ccc", "eee", "aaa"},
+        {true, true, true, false, true, true, true, true, true}}
+        .release();
+
+    auto a_encoded = cudf::dictionary::encode(a->view());
+
+    cudf::transform_input inputs[] = {*a_encoded};
+
+    auto out =
+      cudf::transform_extended(inputs, cuda, cudf::data_type{cudf::type_id::STRING}, false);
+
+    CUDF_TEST_EXPECT_COLUMNS_EQUAL(out->view(), a->view());
+  }
+}
+
+TEST_F(UnaryOperationIntegrationTest, Transform_DictionaryFloat)
+{
+  std::string const cuda =
+    R"***(
+__device__ inline void decode(float * output, float input){
+  *output = input;
+})***";
+
+  // non-nullable
+  {
+    auto a = cudf::test::fixed_width_column_wrapper<float>(
+               {1.0F, 2.0F, 3.0F, 4.0F, 5.0F, 5.0F, 5.0F, 1.0F, 2.0F})
+               .release();
+    auto a_encoded                 = cudf::dictionary::encode(a->view());
+    cudf::transform_input inputs[] = {*a_encoded};
+
+    auto out =
+      cudf::transform_extended(inputs, cuda, cudf::data_type{cudf::type_id::FLOAT32}, false);
+
+    CUDF_TEST_EXPECT_COLUMNS_EQUAL(out->view(), a->view());
+  }
+
+  // nullable
+  {
+    auto a = cudf::test::fixed_width_column_wrapper<float>(
+               {{1.0F, 2.0F, 3.0F, 4.0F, 5.0F, 5.0F, 5.0F, 1.0F, 2.0F},
+                {true, true, true, true, true, false, true, true, true}})
+               .release();
+
+    auto a_encoded                 = cudf::dictionary::encode(a->view());
+    cudf::transform_input inputs[] = {*a_encoded};
+
+    auto out =
+      cudf::transform_extended(inputs, cuda, cudf::data_type{cudf::type_id::FLOAT32}, false);
+
+    CUDF_TEST_EXPECT_COLUMNS_EQUAL(out->view(), a->view());
+  }
+}
+
 struct TernaryOperationTest : public cudf::test::BaseFixture {};
 
 TEST_F(TernaryOperationTest, TransformWithScalar)