[FEA] Unify JIT storage access and mutable string representations

cpp/CMakeLists.txt

@@ -513,7 +513,7 @@ function(precompile_fixed_width_kernel_fragments)
     )
     set(VARIANT_NAME transform_kernel_${FILE_INDEX})
     set(INSTANCE
-        "cudf::jit::transform_kernel<false, false, cudf::jit::type_list<cudf::jit::column_accessor<0, cudf::column_device_view_core, ${TYPE}, false, 0>>, cudf::jit::type_list<cudf::jit::column_accessor<0, cudf::mutable_column_device_view_core, ${TYPE}, false, 0>>>"
+        "cudf::jit::transform_kernel<false, false, cudf::jit::type_list<cudf::jit::column_accessor<0, ${TYPE}, false, 0>>, cudf::jit::type_list<cudf::jit::column_accessor<0, ${TYPE}, false, 0>>>"
     )
     add_fragment(
       cudf_fragments
@@ -548,7 +548,7 @@ function(precompile_fixed_width_kernel_fragments)
       )
       set(VARIANT_NAME transform_kernel_${FILE_INDEX})
       set(INSTANCE
-          "cudf::jit::transform_kernel<false, false, cudf::jit::type_list<cudf::jit::column_accessor<0, cudf::column_device_view_core, ${TYPE}, false, 0>, cudf::jit::column_accessor<1, cudf::column_device_view_core, ${TYPE}, ${RHS_IS_SCALAR}, 0>>, cudf::jit::type_list<cudf::jit::column_accessor<0, cudf::mutable_column_device_view_core, ${TYPE}, false, 0>>>"
+          "cudf::jit::transform_kernel<false, false, cudf::jit::type_list<cudf::jit::column_accessor<0, ${TYPE}, false, 0>, cudf::jit::column_accessor<1, ${TYPE}, ${RHS_IS_SCALAR}, 0>>, cudf::jit::type_list<cudf::jit::column_accessor<0, ${TYPE}, false, 0>>>"
       )
       add_fragment(
         cudf_fragments
@@ -580,18 +580,7 @@ function(precompile_string_kernel_fragments)
   list(APPEND CUDF_PRECOMPILE_STRING_OUTPUT_PHYSICAL_TYPES uint8_t uint16_t uint32_t uint64_t
        cuda::std::span<char>
   )
-  list(
-    APPEND
-    CUDF_PRECOMPILE_STRING_OUTPUT_COLUMN_TYPES
-    cudf::mutable_column_device_view_core
-    cudf::mutable_column_device_view_core
-    cudf::mutable_column_device_view_core
-    cudf::mutable_column_device_view_core
-    cudf::jit::mutable_strings_column_device_view
-  )
-  foreach(OUTPUT_TYPE OUTPUT_COLUMN_TYPE IN ZIP_LISTS CUDF_PRECOMPILE_STRING_OUTPUT_PHYSICAL_TYPES
-                      CUDF_PRECOMPILE_STRING_OUTPUT_COLUMN_TYPES
-  )
+  foreach(OUTPUT_TYPE IN LISTS CUDF_PRECOMPILE_STRING_OUTPUT_PHYSICAL_TYPES)
     # Pre-compile for up to 8 output columns
     foreach(OUTPUT_INDICES IN ITEMS "0" "0;1" "0;1;2" "0;1;2;3" "0;1;2;3;4" "0;1;2;3;4;5"
                                     "0;1;2;3;4;5;6" "0;1;2;3;4;5;6;7"
@@ -606,17 +595,15 @@ function(precompile_string_kernel_fragments)
       set(OUTPUT_ACCESSORS "")
       set(OUTPUT_POINTERS "")
       foreach(OUTPUT_INDEX IN LISTS OUTPUT_INDICES)
-        list(
-          APPEND
-          OUTPUT_ACCESSORS
-          "cudf::jit::column_accessor<${OUTPUT_INDEX}, ${OUTPUT_COLUMN_TYPE}, ${OUTPUT_TYPE}, false, 0>"
+        list(APPEND OUTPUT_ACCESSORS
+             "cudf::jit::column_accessor<${OUTPUT_INDEX}, ${OUTPUT_TYPE}, false, 0>"
         )
         list(APPEND OUTPUT_POINTERS "${OUTPUT_TYPE} *")
       endforeach()
       list(JOIN OUTPUT_ACCESSORS " ," OUTPUT_ACCESSORS_STR)
       list(JOIN OUTPUT_POINTERS " ," OUTPUT_POINTERS_STR)
       set(INSTANCE
-          "cudf::jit::transform_kernel<false, false, cudf::jit::type_list<cudf::jit::column_accessor<0, cudf::column_device_view_core, cudf::string_view, false, 0>>, cudf::jit::type_list<${OUTPUT_ACCESSORS_STR}>>"
+          "cudf::jit::transform_kernel<false, false, cudf::jit::type_list<cudf::jit::column_accessor<0, cudf::string_view, false, 0>>, cudf::jit::type_list<${OUTPUT_ACCESSORS_STR}>>"
       )
       add_fragment(
         cudf_fragments

cpp/include/cudf/column/column_device_view_base.cuh

@@ -14,6 +14,7 @@
 
 #include <cuda/std/algorithm>
 #include <cuda/std/optional>
+#include <cuda/std/span>
 #include <cuda/std/type_traits>
 
 /**
@@ -92,6 +93,33 @@ inline constexpr bool is_dictionary_encoded = false;
 template <typename IndexType, typename KeyType>
 inline constexpr bool is_dictionary_encoded<dictionary_element<IndexType, KeyType>> = true;
 
+/**
+ * @brief An element interpreted directly as its byte representation.
+ * @tparam T The underlying type used to store the value (e.g., `bool` or `uint8_t`, `char`,
+ * `string_view`, etc.)
+ *
+ */
+template <typename T>
+struct repr_element {
+  using value_type = T;  ///< The underlying type used to store the value
+
+  T value{};  ///< The value for this element
+};
+
+/**
+ * @brief A type trait to determine if a type is a representation element.
+ * @tparam T The type to check
+ */
+template <typename T>
+inline constexpr bool is_repr_element = false;
+
+/**
+ * @brief Specialization of `is_repr_element` for `repr_element<T>`.
+ * @tparam T The underlying type used to store the value
+ */
+template <typename T>
+inline constexpr bool is_repr_element<repr_element<T>> = true;
+
 namespace detail {
 /**
  * @brief An immutable, non-owning view of device data as a column of elements
@@ -466,7 +494,8 @@ class alignas(16) column_device_view_core : public detail::column_device_view_ba
    * @param element_index Position of the desired string element
    * @return string_view instance representing this element at this index
    */
-  template <typename T, CUDF_ENABLE_IF(cuda::std::is_same_v<T, string_view>)>
+  template <typename T,
+            CUDF_ENABLE_IF(cuda::std::is_same_v<cuda::std::remove_cv_t<T>, string_view>)>
   [[nodiscard]] __device__ T element(size_type element_index) const noexcept
   {
     size_type index       = element_index + offset();  // account for this view's _offset
@@ -515,7 +544,7 @@ class alignas(16) column_device_view_core : public detail::column_device_view_ba
    * @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
+  [[nodiscard]] __device__ auto element(size_type element_index) const noexcept
   {
     auto const& indices = child(dictionary_indices_column_index);
     auto const& keys    = child(dictionary_keys_column_index);
@@ -524,20 +553,34 @@ class alignas(16) column_device_view_core : public detail::column_device_view_ba
     return keys.template element<typename T::key_type>(index);
   }
 
+  /**
+   * @brief Reads a raw child value.
+   *
+   * The child must be valid. Use nullable_element<T>() to check child validity.
+   * @tparam T The raw child type, satisfying `is_repr_element<T>`
+   * @param element_index Position of the desired element in `[0, size())`
+   * @return The child value
+   */
+  template <typename T, CUDF_ENABLE_IF(is_repr_element<T>)>
+  [[nodiscard]] __device__ auto element(size_type element_index) const noexcept
+  {
+    return static_cast<typename T::value_type const*>(_data)[_offset + element_index];
+  }
+
   /**
    * @brief Returns a nullable element at the specified index. If the element is null, returns
    * `nullopt`.
    *
    * @param element_index Position of the desired element
-   * @return `optional` containing the element at the specified index, or `nullopt` if the element
-   * is null
+   * @return `optional` containing the decoded value returned by `element<T>()`, or `nullopt`
+   * if the element is null. For dictionary tags, the optional contains the decoded key type.
    */
   template <typename T>
-  [[nodiscard]] __device__ cuda::std::optional<T> nullable_element(
-    size_type element_index) const noexcept
+  [[nodiscard]] __device__ auto nullable_element(size_type element_index) const noexcept
   {
-    if (is_null(element_index)) { return cuda::std::nullopt; }
-    return element<T>(element_index);
+    using value_type = cuda::std::remove_cvref_t<decltype(element<T>(element_index))>;
+    if (is_null(element_index)) { return cuda::std::optional<value_type>{}; }
+    return cuda::std::optional<value_type>{element<T>(element_index)};
   }
 
   /**
@@ -699,7 +742,7 @@ class alignas(16) mutable_column_device_view_core : public detail::column_device
    * @param element_index Position of the desired string element
    * @return string_view instance representing this element at this index
    */
-  template <typename T, CUDF_ENABLE_IF(cuda::std::is_same_v<T, string_view>)>
+  template <typename T, CUDF_ENABLE_IF(cuda::std::same_as<cuda::std::remove_cv_t<T>, string_view>)>
   [[nodiscard]] __device__ T element(size_type element_index) const noexcept
   {
     size_type index       = element_index + offset();  // account for this view's _offset
@@ -729,20 +772,55 @@ class alignas(16) mutable_column_device_view_core : public detail::column_device
     return T{scaled_integer<rep>{data<rep>()[element_index], scale}};
   }
 
+  /**
+   * @brief Returns a mutable view into the existing string character buffer.
+   *
+   * The column offset is applied. The element must be valid, and mutations must stay
+   * within the existing string's byte range.
+   * @tparam T `cuda::std::span<char>`
+   * @param element_index Position of the desired string element
+   * @return A mutable view of the string's preallocated character storage
+   */
+  template <typename T,
+            CUDF_ENABLE_IF(cuda::std::same_as<cuda::std::remove_cv_t<T>, cuda::std::span<char>>)>
+  [[nodiscard]] __device__ T element(size_type element_index) const noexcept
+  {
+    auto const index   = element_index + offset();
+    auto const offsets = child(offsets_column_index);
+    auto const itr     = detail::input_offsetalator(offsets.head(), offsets.type());
+    auto const begin   = itr[index];
+    auto const end     = itr[index + 1];
+    return T{head<char>() + begin, static_cast<size_t>(end - begin)};
+  }
+
+  /**
+   * @brief Reads a raw child value.
+   *
+   * The child must be valid. Use nullable_element<T>() to check child validity.
+   * @tparam T The raw child type, satisfying `is_repr_element<T>`
+   * @param element_index Position of the desired element in `[0, size())`
+   * @return The child value
+   */
+  template <typename T, CUDF_ENABLE_IF(is_repr_element<T>)>
+  [[nodiscard]] __device__ auto element(size_type element_index) const noexcept
+  {
+    return static_cast<typename T::value_type*>(const_cast<void*>(_data))[_offset + element_index];
+  }
+
   /**
    * @brief Returns a nullable element at the specified index. If the element is null, returns
    * `nullopt`.
    *
    * @param element_index Position of the desired element
-   * @return `optional` containing the element at the specified index, or `nullopt` if the element
-   * is null
+   * @return `optional` containing the decoded value returned by `element<T>()`, or `nullopt`
+   * if the element is null. For dictionary tags, the optional contains the decoded key type.
    */
   template <typename T>
-  [[nodiscard]] __device__ cuda::std::optional<T> nullable_element(
-    size_type element_index) const noexcept
+  [[nodiscard]] __device__ auto nullable_element(size_type element_index) const noexcept
   {
-    if (is_null(element_index)) { return cuda::std::nullopt; }
-    return element<T>(element_index);
+    using value_type = cuda::std::remove_cvref_t<decltype(element<T>(element_index))>;
+    if (is_null(element_index)) { return cuda::std::optional<value_type>{}; }
+    return cuda::std::optional<value_type>{element<T>(element_index)};
   }
 
   /**
@@ -775,6 +853,34 @@ class alignas(16) mutable_column_device_view_core : public detail::column_device
     data<rep>()[element_index] = value.value();
   }
 
+  /**
+   * @brief Assigns `value` to the element at `element_index`.
+   * This is a no-op for mutable string elements.
+   *
+   * @tparam T The element type
+   * @param element_index Position of the desired element
+   * @param value The value to assign
+   */
+  template <typename T,
+            CUDF_ENABLE_IF(cuda::std::same_as<cuda::std::remove_cv_t<T>, cuda::std::span<char>>)>
+  __device__ void assign(size_type element_index, T value) const noexcept
+  {
+    // no-op for mutable string elements
+  }
+
+  /**
+   * @brief Assigns `value` to the element at `element_index`
+   *
+   * @tparam T The element type
+   * @param element_index Position of the desired element
+   * @param value The value to assign
+   */
+  template <typename T, CUDF_ENABLE_IF(is_repr_element<T>)>
+  __device__ void assign(size_type element_index, typename T::value_type value) const noexcept
+  {
+    static_cast<typename T::value_type*>(const_cast<void*>(_data))[_offset + element_index] = value;
+  }
+
   /**
    * @brief Returns raw pointer to the underlying bitmask allocation.
    *

cpp/src/jit/column_accessor.cuh

@@ -8,25 +8,20 @@
 #include <cudf/column/column_device_view_base.cuh>
 #include <cudf/types.hpp>
 
-#include <cuda/std/cstddef>
 #include <cuda/std/optional>
-#include <cuda/std/span>
 #include <cuda/std/type_traits>
 
 namespace cudf {
 namespace jit {
 
-template <int32_t Index, typename Column, typename Element, bool AsScalar, int32_t TableIndex>
+template <int32_t Index, typename Element, bool AsScalar, int32_t TableIndex>
 struct column_accessor {
   static constexpr int32_t index       = Index;
-  static constexpr int32_t table_index = TableIndex;
-  using column_type                    = Column;
   using element_type                   = Element;
-  using optional_element_type          = cuda::std::optional<element_type>;
-
-  static constexpr bool as_scalar = AsScalar;
+  static constexpr bool as_scalar      = AsScalar;
+  static constexpr int32_t table_index = TableIndex;
 
-  static constexpr bool is_mutable_string =
+  static constexpr bool is_inplace_output =
     cuda::std::is_same_v<element_type, cuda::std::span<char>>;
 
   static __device__ constexpr size_type map_index(size_type row)
@@ -38,69 +33,67 @@ struct column_accessor {
     }
   }
 
-  template <typename T>
-  static __device__ auto& column(T const* __restrict__ cols)
-    requires(sizeof(T) == sizeof(column_type) && alignof(T) == alignof(column_type))
-  {
-    return reinterpret_cast<column_type const&>(cols[index]);
-  }
+  static __device__ auto& column(auto const* __restrict__ columns) { return columns[index]; }
 
-  static __device__ auto element(auto const* __restrict__ cols, size_type row)
+  static __device__ auto element(auto const* __restrict__ columns, size_type row)
   {
-    return column(cols).template element<element_type>(map_index(row));
+    return columns[index].template element<element_type>(map_index(row));
   }
 
-  static __device__ bool is_null(auto const* __restrict__ cols, size_type row)
+  static __device__ bool is_null(auto const* __restrict__ columns, size_type row)
   {
-    return column(cols).is_null(map_index(row));
+    return columns[index].is_null(map_index(row));
   }
 
-  static __device__ bool is_valid(auto const* __restrict__ cols, size_type row)
+  static __device__ bool is_valid(auto const* __restrict__ columns, size_type row)
   {
-    return column(cols).is_valid(map_index(row));
+    return columns[index].is_valid(map_index(row));
   }
 
-  static __device__ auto nullable_element(auto const* __restrict__ cols, size_type row)
+  static __device__ auto nullable_element(auto const* __restrict__ columns, size_type row)
   {
-    return column(cols).template nullable_element<element_type>(map_index(row));
+    return columns[index].template nullable_element<element_type>(map_index(row));
   }
 
-  static __device__ void set_null_mask_word(auto const* __restrict__ cols,
+  static __device__ void set_null_mask_word(auto const* __restrict__ columns,
                                             size_type word_index,
                                             bitmask_type word)
     requires(!as_scalar)
   {
-    auto* mask = column(cols).null_mask();
+    auto* mask = columns[index].null_mask();
 
     if (mask == nullptr) { return; }
 
     mask[word_index] = word;
   }
 
-  static __device__ void assign(auto const* __restrict__ cols, size_type row, element_type value)
+  static __device__ void assign(auto const* __restrict__ columns, size_type row, auto value)
     requires(!as_scalar)
   {
-    column(cols).template assign<element_type>(row, value);
+    columns[index].template assign<element_type>(row, value);
   }
 
-  static __device__ element_type output_arg(auto const* __restrict__ cols, size_type row)
+  static __device__ auto output_arg(auto const* __restrict__ columns, size_type row)
     requires(!as_scalar)
   {
-    if constexpr (is_mutable_string) {
-      return element(cols, row);
+    if constexpr (is_inplace_output) {
+      return element(columns, row);
+    } else if constexpr (is_repr_element<element_type>) {
+      return typename element_type::value_type{};
     } else {
-      return {};
+      return element_type{};
     }
   }
 
-  static __device__ optional_element_type null_output_arg(auto const* __restrict__ cols,
-                                                          size_type row)
+  static __device__ auto nullable_output_arg(auto const* __restrict__ columns, size_type row)
     requires(!as_scalar)
   {
-    if constexpr (is_mutable_string) {
-      return element(cols, row);
+    if constexpr (is_inplace_output) {
+      return cuda::std::optional<element_type>(element(columns, row));
+    } else if constexpr (is_repr_element<element_type>) {
+      return cuda::std::optional<typename element_type::value_type>{};
     } else {
-      return {};
+      return cuda::std::optional<element_type>{};
     }
   }
 };