diff --git a/src/data/array_interface.h b/src/data/array_interface.h index 1c07f2bb0..53152b63d 100644 --- a/src/data/array_interface.h +++ b/src/data/array_interface.h @@ -231,7 +231,8 @@ class ArrayInterfaceHandler { class ArrayInterface { public: ArrayInterface() = default; - explicit ArrayInterface(std::map const& column) { + explicit ArrayInterface(std::map const &column, + bool allow_mask = true) { ArrayInterfaceHandler::Validate(column); data = ArrayInterfaceHandler::GetPtrFromArrayData(column); CHECK(data) << "Column is null"; @@ -239,16 +240,22 @@ class ArrayInterface { num_rows = shape.first; num_cols = shape.second; - common::Span s_mask; - size_t n_bits = ArrayInterfaceHandler::ExtractMask(column, &s_mask); + if (allow_mask) { + common::Span s_mask; + size_t n_bits = ArrayInterfaceHandler::ExtractMask(column, &s_mask); - valid = RBitField8(s_mask); + valid = RBitField8(s_mask); - if (s_mask.data()) { - CHECK_EQ(n_bits, num_rows) - << "Shape of bit mask doesn't match data shape. " - << "XGBoost doesn't support internal broadcasting."; + if (s_mask.data()) { + CHECK_EQ(n_bits, num_rows) + << "Shape of bit mask doesn't match data shape. " + << "XGBoost doesn't support internal broadcasting."; + } + } else { + CHECK(column.find("mask") == column.cend()) + << "Masked array is not yet supported."; } + auto typestr = get(column.at("typestr")); type[0] = typestr.at(0); type[1] = typestr.at(1); diff --git a/src/data/device_adapter.cuh b/src/data/device_adapter.cuh index f52b45aa1..d2e039e48 100644 --- a/src/data/device_adapter.cuh +++ b/src/data/device_adapter.cuh @@ -177,10 +177,7 @@ class CupyAdapterBatch : public detail::NoMetaInfo { __device__ COOTuple GetElement(size_t idx) const { size_t column_idx = idx % array_interface_.num_cols; size_t row_idx = idx / array_interface_.num_cols; - float value = array_interface_.valid.Data() == nullptr || - array_interface_.valid.Check(row_idx) - ? array_interface_.GetElement(idx) - : std::numeric_limits::quiet_NaN(); + float value = array_interface_.GetElement(idx); return {row_idx, column_idx, value}; } @@ -193,7 +190,7 @@ class CupyAdapter : public detail::SingleBatchDataIter { explicit CupyAdapter(std::string cuda_interface_str) { Json json_array_interface = Json::Load({cuda_interface_str.c_str(), cuda_interface_str.size()}); - array_interface_ = ArrayInterface(get(json_array_interface)); + array_interface_ = ArrayInterface(get(json_array_interface), false); device_idx_ = dh::CudaGetPointerDevice(array_interface_.data); CHECK_NE(device_idx_, -1); batch_ = CupyAdapterBatch(array_interface_);