Store categories from pandas. (#11303)

- Change host columnar adapter to receive categories.
- Store categories in the meta info.
- Glue code for the Python interface.

Author: trivialfis

R-package/src/Makevars.in

@@ -66,6 +66,7 @@ OBJECTS= \
     $(PKGROOT)/src/gbm/gblinear_model.o \
     $(PKGROOT)/src/data/adapter.o \
     $(PKGROOT)/src/data/array_interface.o \
+    $(PKGROOT)/src/data/cat_container.o \
     $(PKGROOT)/src/data/simple_dmatrix.o \
     $(PKGROOT)/src/data/data.o \
     $(PKGROOT)/src/data/sparse_page_raw_format.o \

R-package/src/Makevars.win.in

@@ -65,6 +65,7 @@ OBJECTS= \
     $(PKGROOT)/src/gbm/gblinear_model.o \
     $(PKGROOT)/src/data/adapter.o \
     $(PKGROOT)/src/data/array_interface.o \
+    $(PKGROOT)/src/data/cat_container.o \
     $(PKGROOT)/src/data/simple_dmatrix.o \
     $(PKGROOT)/src/data/data.o \
     $(PKGROOT)/src/data/sparse_page_raw_format.o \

include/xgboost/data.h

@@ -1,5 +1,5 @@
 /**
- * Copyright 2015-2024, XGBoost Contributors
+ * Copyright 2015-2025, XGBoost Contributors
  * \file data.h
  * \brief The input data structure of xgboost.
  * \author Tianqi Chen
@@ -8,8 +8,8 @@
 #define XGBOOST_DATA_H_
 
 #include <dmlc/base.h>
-#include <dmlc/data.h>
-#include <dmlc/serializer.h>
+#include <dmlc/io.h>          // for Stream
+#include <dmlc/serializer.h>  // for Handler
 #include <xgboost/base.h>
 #include <xgboost/host_device_vector.h>
 #include <xgboost/linalg.h>
@@ -42,13 +42,16 @@ enum class FeatureType : uint8_t { kNumerical = 0, kCategorical = 1 };
 
 enum class DataSplitMode : int { kRow = 0, kCol = 1 };
 
-/*!
- * \brief Meta information about dataset, always sit in memory.
+// Forward declaration of the container used by the meta info.
+struct CatContainer;
+
+/**
+ * @brief Meta information about dataset, always sit in memory.
  */
 class MetaInfo {
  public:
   /*! \brief number of data fields in MetaInfo */
-  static constexpr uint64_t kNumField = 12;
+  static constexpr uint64_t kNumField = 13;
 
   /*! \brief number of rows in the data */
   bst_idx_t num_row_{0};  // NOLINT
@@ -100,9 +103,9 @@ class MetaInfo {
    */
   HostDeviceVector<float> feature_weights;
 
-  /*! \brief default constructor */
-  MetaInfo()  = default;
+  MetaInfo();
   MetaInfo(MetaInfo&& that) = default;
+  MetaInfo(MetaInfo const& that) = delete;
   MetaInfo& operator=(MetaInfo&& that) = default;
   MetaInfo& operator=(MetaInfo const& that) = delete;
 
@@ -205,6 +208,16 @@ class MetaInfo {
    * @brief Flag for whether the DMatrix has categorical features.
    */
   bool HasCategorical() const { return has_categorical_; }
+  /**
+   * @brief Getters for categories.
+   */
+  [[nodiscard]] CatContainer const* Cats() const;
+  [[nodiscard]] CatContainer* Cats();
+  [[nodiscard]] std::shared_ptr<CatContainer const>  CatsShared() const;
+  /**
+   * @brief Setter for categories.
+   */
+  void Cats(std::shared_ptr<CatContainer> cats);
 
  private:
   void SetInfoFromHost(Context const* ctx, StringView key, Json arr);
@@ -213,6 +226,8 @@ class MetaInfo {
   /*! \brief argsort of labels */
   mutable std::vector<size_t> label_order_cache_;
   bool has_categorical_{false};
+
+  std::shared_ptr<CatContainer> cats_;
 };
 
 /*! \brief Element from a sparse vector */
@@ -691,7 +706,15 @@ class DMatrix {
    * @param slice_id Index of the current slice
    * @return DMatrix containing the slice of columns
    */
-  virtual DMatrix *SliceCol(int num_slices, int slice_id) = 0;
+  virtual DMatrix* SliceCol(int num_slices, int slice_id) = 0;
+  /**
+   * @brief Accessor for the string representation of the categories.
+   */
+  CatContainer const* Cats() const { return this->CatsShared().get(); }
+  [[nodiscard]] virtual std::shared_ptr<CatContainer const> CatsShared() const {
+    LOG(FATAL) << "Not implemented for the current DMatrix type.";
+    return nullptr;
+  }
 
  protected:
   virtual BatchSet<SparsePage> GetRowBatches() = 0;

include/xgboost/predictor.h

@@ -5,7 +5,8 @@
  *  performs predictions for a gradient booster.
  */
 #pragma once
-#include <xgboost/base.h>
+#include <dmlc/registry.h>    // for FunctionRegEntryBase
+#include <xgboost/base.h>     // for bst_tree_t
 #include <xgboost/cache.h>    // for DMatrixCache
 #include <xgboost/context.h>  // for Context
 #include <xgboost/context.h>

ops/script/lint_python.py

@@ -27,6 +27,7 @@ class LintersPaths:
         "tests/python/test_early_stopping.py",
         "tests/python/test_multi_target.py",
         "tests/python/test_objectives.py",
+        "tests/python/test_ordinal.py",
         "tests/python/test_predict.py",
         "tests/python/test_quantile_dmatrix.py",
         "tests/python/test_tracker.py",
@@ -101,6 +102,7 @@ class LintersPaths:
         "tests/python-gpu/load_pickle.py",
         "tests/python-gpu/test_gpu_training_continuation.py",
         "tests/python/test_model_io.py",
+        "tests/python/test_ordinal.py",
         "tests/test_distributed/test_federated/",
         "tests/test_distributed/test_gpu_federated/",
         "tests/test_distributed/test_with_dask/test_ranking.py",

python-package/xgboost/_data_utils.py

@@ -2,15 +2,35 @@
 
 import copy
 import ctypes
+import functools
 import json
-from typing import Literal, Optional, Protocol, Tuple, Type, TypedDict, Union, cast
+from typing import (
+    TYPE_CHECKING,
+    Any,
+    Dict,
+    Literal,
+    Optional,
+    Protocol,
+    Tuple,
+    Type,
+    TypedDict,
+    TypeGuard,
+    Union,
+    cast,
+    overload,
+)
 
 import numpy as np
 
-from ._typing import CNumericPtr, DataType, NumpyOrCupy
-from .compat import import_cupy
+from ._typing import CNumericPtr, DataType, NumpyDType, NumpyOrCupy
+from .compat import import_cupy, lazy_isinstance
+
+if TYPE_CHECKING:
+    import pandas as pd
+    import pyarrow as pa
 
 
+# Used for accepting inputs for numpy and cupy arrays
 class _ArrayLikeArg(Protocol):
     @property
     def __array_interface__(self) -> "ArrayInf": ...
@@ -44,19 +64,27 @@ def shape(self) -> Tuple[int, int]:
     },
 )
 
+StringArray = TypedDict("StringArray", {"offsets": ArrayInf, "values": ArrayInf})
+
 
 def array_hasobject(data: DataType) -> bool:
     """Whether the numpy array has object dtype."""
     return hasattr(data.dtype, "hasobject") and data.dtype.hasobject
 
 
-def cuda_array_interface(data: DataType) -> bytes:
-    """Make cuda array interface str."""
+def cuda_array_interface_dict(data: _CudaArrayLikeArg) -> ArrayInf:
+    """Returns a dictionary storing the CUDA array interface."""
     if array_hasobject(data):
         raise ValueError("Input data contains `object` dtype.  Expecting numeric data.")
-    interface = data.__cuda_array_interface__
-    if "mask" in interface:
-        interface["mask"] = interface["mask"].__cuda_array_interface__
+    ainf = data.__cuda_array_interface__
+    if "mask" in ainf:
+        ainf["mask"] = ainf["mask"].__cuda_array_interface__  # type: ignore
+    return cast(ArrayInf, ainf)
+
+
+def cuda_array_interface(data: _CudaArrayLikeArg) -> bytes:
+    """Make cuda array interface str."""
+    interface = cuda_array_interface_dict(data)
     interface_str = bytes(json.dumps(interface), "utf-8")
     return interface_str
 
@@ -107,6 +135,12 @@ def __cuda_array_interface__(self, interface: ArrayInf) -> None:
     return out
 
 
+# Default constant value for CUDA per-thread stream.
+STREAM_PER_THREAD = 2
+
+
+# Typing is not strict as there are subtle differences between CUDA array interface and
+# array interface. We handle them uniformly for now.
 def make_array_interface(
     ptr: Union[CNumericPtr, int],
     shape: Tuple[int, ...],
@@ -134,21 +168,157 @@ def make_array_interface(
         return array
 
     array["data"] = (addr, True)
-    if is_cuda:
-        array["stream"] = 2
+    if is_cuda and "stream" not in array:
+        array["stream"] = STREAM_PER_THREAD
     array["shape"] = shape
     array["strides"] = None
     return array
 
 
-def array_interface_dict(data: np.ndarray) -> ArrayInf:
-    """Convert array interface into a Python dictionary."""
+def is_arrow_dict(data: Any) -> TypeGuard["pa.DictionaryArray"]:
+    """Is this an arrow dictionary array?"""
+    return lazy_isinstance(data, "pyarrow.lib", "DictionaryArray")
+
+
+class PdCatAccessor(Protocol):
+    """Protocol for pandas cat accessor."""
+
+    @property
+    def categories(  # pylint: disable=missing-function-docstring
+        self,
+    ) -> "pd.Index": ...
+
+    @property
+    def codes(self) -> "pd.Series": ...  # pylint: disable=missing-function-docstring
+
+    @property
+    def dtype(self) -> np.dtype: ...  # pylint: disable=missing-function-docstring
+
+    def to_arrow(  # pylint: disable=missing-function-docstring
+        self,
+    ) -> Union["pa.StringArray", "pa.IntegerArray"]: ...
+
+    @property
+    def __cuda_array_interface__(self) -> ArrayInf: ...
+
+
+def _is_pd_cat(data: Any) -> TypeGuard[PdCatAccessor]:
+    # Test pd.Series.cat, not pd.Series
+    return hasattr(data, "categories") and hasattr(data, "codes")
+
+
+@functools.cache
+def _arrow_typestr() -> Dict["pa.DataType", str]:
+    import pyarrow as pa
+
+    mapping = {
+        pa.int8(): "<i1",
+        pa.int16(): "<i2",
+        pa.int32(): "<i4",
+        pa.int64(): "<i8",
+        pa.uint8(): "<u1",
+        pa.uint16(): "<u2",
+        pa.uint32(): "<u4",
+        pa.uint64(): "<u8",
+    }
+
+    return mapping
+
+
+def npstr_to_arrow_strarr(strarr: np.ndarray) -> Tuple[np.ndarray, str]:
+    """Convert a numpy string array to an arrow string array."""
+    lenarr = np.vectorize(len)
+    offsets = np.cumsum(np.concatenate([np.array([0], dtype=np.int64), lenarr(strarr)]))
+    values = strarr.sum()
+    assert "\0" not in values  # arrow string array doesn't need null terminal
+    return offsets.astype(np.int32), values
+
+
+def _ensure_np_dtype(
+    data: DataType, dtype: Optional[NumpyDType]
+) -> Tuple[np.ndarray, Optional[NumpyDType]]:
+    """Ensure the np array has correct type and is contiguous."""
+    if array_hasobject(data) or data.dtype in [np.float16, np.bool_]:
+        dtype = np.float32
+        data = data.astype(dtype, copy=False)
+    if not data.flags.aligned:
+        data = np.require(data, requirements="A")
+    return data, dtype
+
+
+@overload
+def array_interface_dict(data: np.ndarray) -> ArrayInf: ...
+
+
+@overload
+def array_interface_dict(
+    data: PdCatAccessor,
+) -> Tuple[StringArray, ArrayInf, Tuple]: ...
+
+
+@overload
+def array_interface_dict(
+    data: "pa.DictionaryArray",
+) -> Tuple[StringArray, ArrayInf, Tuple]: ...
+
+
+def array_interface_dict(  # pylint: disable=too-many-locals
+    data: Union[np.ndarray, PdCatAccessor],
+) -> Union[ArrayInf, Tuple[StringArray, ArrayInf, Optional[Tuple]]]:
+    """Returns an array interface from the input."""
+    # Handle categorical values
+    if _is_pd_cat(data):
+        cats = data.categories
+        # pandas uses -1 to represent missing values for categorical features
+        codes = data.codes.replace(-1, np.nan)
+
+        if np.issubdtype(cats.dtype, np.floating) or np.issubdtype(
+            cats.dtype, np.integer
+        ):
+            # Numeric index type
+            name_values = cats.values
+            jarr_values = array_interface_dict(name_values)
+            code_values = codes.values
+            jarr_codes = array_interface_dict(code_values)
+            return jarr_values, jarr_codes, (name_values, code_values)
+
+        # String index type
+        name_offsets, name_values = npstr_to_arrow_strarr(cats.values)
+        name_offsets, _ = _ensure_np_dtype(name_offsets, np.int32)
+        joffsets = array_interface_dict(name_offsets)
+        bvalues = name_values.encode("utf-8")
+        ptr = ctypes.c_void_p.from_buffer(ctypes.c_char_p(bvalues)).value
+        assert ptr is not None
+
+        jvalues: ArrayInf = {
+            "data": (ptr, True),
+            "typestr": "|i1",
+            "shape": (len(name_values),),
+            "strides": None,
+            "version": 3,
+            "mask": None,
+        }
+        jnames: StringArray = {"offsets": joffsets, "values": jvalues}
+
+        code_values = codes.values
+        jcodes = array_interface_dict(code_values)
+
+        buf = (
+            name_offsets,
+            name_values,
+            bvalues,
+            code_values,
+        )  # store temporary values
+        return jnames, jcodes, buf
+
+    # Handle numeric values
+    assert isinstance(data, np.ndarray)
     if array_hasobject(data):
         raise ValueError("Input data contains `object` dtype.  Expecting numeric data.")
-    arrinf = data.__array_interface__
-    if "mask" in arrinf:
-        arrinf["mask"] = arrinf["mask"].__array_interface__
-    return cast(ArrayInf, arrinf)
+    ainf = data.__array_interface__
+    if "mask" in ainf:
+        ainf["mask"] = ainf["mask"].__array_interface__
+    return cast(ArrayInf, ainf)
 
 
 def array_interface(data: np.ndarray) -> bytes: