Fix CSR Array adapter batch columns.

src/data/adapter.h

@@ -288,6 +288,7 @@ class CSRArrayAdapterBatch : public detail::NoMetaInfo {
   ArrayInterface indptr_;
   ArrayInterface indices_;
   ArrayInterface values_;
+  bst_feature_t n_features_;
 
   class Line {
     ArrayInterface indices_;
@@ -311,23 +312,27 @@ class CSRArrayAdapterBatch : public detail::NoMetaInfo {
     }
   };
 
+ public:
+  static constexpr bool kIsRowMajor = true;
+
  public:
   CSRArrayAdapterBatch() = default;
   CSRArrayAdapterBatch(ArrayInterface indptr, ArrayInterface indices,
-                       ArrayInterface values)
+                       ArrayInterface values, bst_feature_t n_features)
       : indptr_{std::move(indptr)}, indices_{std::move(indices)},
-        values_{std::move(values)} {
+        values_{std::move(values)}, n_features_{n_features} {
     indptr_.AsColumnVector();
     values_.AsColumnVector();
     indices_.AsColumnVector();
   }
 
-  size_t Size() const {
+  size_t NumRows() const {
     size_t size = indptr_.num_rows * indptr_.num_cols;
     size = size == 0 ? 0 : size - 1;
     return size;
   }
-  static constexpr bool kIsRowMajor = true;
+  size_t NumCols() const { return n_features_; }
+  size_t Size() const { return this->NumRows(); }
 
   Line const GetLine(size_t idx) const {
     auto begin_offset = indptr_.GetElement<size_t>(idx, 0);
@@ -356,7 +361,8 @@ class CSRArrayAdapter : public detail::SingleBatchDataIter<CSRArrayAdapterBatch>
   CSRArrayAdapter(StringView indptr, StringView indices, StringView values,
                   size_t num_cols)
       : indptr_{indptr}, indices_{indices}, values_{values}, num_cols_{num_cols} {
-    batch_ = CSRArrayAdapterBatch{indptr_, indices_, values_};
+    batch_ = CSRArrayAdapterBatch{indptr_, indices_, values_,
+                                  static_cast<bst_feature_t>(num_cols_)};
   }
 
   CSRArrayAdapterBatch const& Value() const override {

src/data/array_interface.h

@@ -11,6 +11,7 @@
 #include <map>
 #include <string>
 #include <utility>
+#include <vector>
 
 #include "xgboost/base.h"
 #include "xgboost/data.h"
@@ -416,5 +417,24 @@ class ArrayInterface {
   Type type;
 };
 
+template <typename T> std::string MakeArrayInterface(T const *data, size_t n) {
+  Json arr{Object{}};
+  arr["data"] = Array(std::vector<Json>{
+      Json{Integer{reinterpret_cast<int64_t>(data)}}, Json{Boolean{false}}});
+  arr["shape"] = Array{std::vector<Json>{Json{Integer{n}}, Json{Integer{1}}}};
+  std::string typestr;
+  if (DMLC_LITTLE_ENDIAN) {
+    typestr.push_back('<');
+  } else {
+    typestr.push_back('>');
+  }
+  typestr.push_back(ArrayInterfaceHandler::TypeChar<T>());
+  typestr += std::to_string(sizeof(T));
+  arr["typestr"] = typestr;
+  arr["version"] = 3;
+  std::string str;
+  Json::Dump(arr, &str);
+  return str;
+}
 }  // namespace xgboost
 #endif  // XGBOOST_DATA_ARRAY_INTERFACE_H_

tests/cpp/data/test_adapter.cc

@@ -35,6 +35,27 @@ TEST(Adapter, CSRAdapter) {
   EXPECT_EQ(line2.GetElement(0).column_idx, 1);
 }
 
+TEST(Adapter, CSRArrayAdapter) {
+  HostDeviceVector<bst_row_t> indptr;
+  HostDeviceVector<float> values;
+  HostDeviceVector<bst_feature_t> indices;
+  size_t n_features = 100, n_samples = 10;
+  RandomDataGenerator{n_samples, n_features, 0.5}.GenerateCSR(&values, &indptr, &indices);
+  auto indptr_arr = MakeArrayInterface(indptr.HostPointer(), indptr.Size());
+  auto values_arr = MakeArrayInterface(values.HostPointer(), values.Size());
+  auto indices_arr = MakeArrayInterface(indices.HostPointer(), indices.Size());
+  auto adapter = data::CSRArrayAdapter(
+      StringView{indptr_arr.c_str(), indptr_arr.size()},
+      StringView{values_arr.c_str(), values_arr.size()},
+      StringView{indices_arr.c_str(), indices_arr.size()}, n_features);
+  auto batch = adapter.Value();
+  ASSERT_EQ(batch.NumRows(), n_samples);
+  ASSERT_EQ(batch.NumCols(), n_features);
+
+  ASSERT_EQ(adapter.NumRows(), n_samples);
+  ASSERT_EQ(adapter.NumColumns(), n_features);
+}
+
 TEST(Adapter, CSCAdapterColsMoreThanRows) {
   std::vector<float> data = {1, 2, 3, 4, 5, 6, 7, 8};
   std::vector<unsigned> row_idx = {0, 1, 0, 1, 0, 1, 0, 1};