Skip to content

Commit

Permalink
decimal128 Support for to/from_arrow (#9986)
Browse files Browse the repository at this point in the history 
Resolves C++ side of #9980.

The reason this PR is breaking is because Arrow only has a notion of `decimal128` (see `arrow::Type::DECIMAL`). We can still support both `decimal64` **and** `decimal128` for `to_arrow` but for `from_arrow` it only makes sense to support one of them, and `decimal128` (now that we have it) is the logical choice. Therfore, the switching of the return type of a column coming `from_arrow` from `decimal64` to `decimal128` is a breaking change.

Requires:
* #7314
* #9533

Authors:
   - Conor Hoekstra (https://github.com/codereport)

Approvers:
   - Devavret Makkar (https://github.com/devavret)
   - Mike Wilson (https://github.com/hyperbolic2346)
  • Loading branch information
@codereport
codereport authored Jan 18, 2022
1 parent 5ea3df6 commit 45c20d1
Show file tree
Hide file tree
Showing 5 changed files with 218 additions and 121 deletions.
31 changes: 12 additions & 19 deletions cpp/src/interop/from_arrow.cu
View file
Original file line number Diff line number Diff line change
@@ -1,5 +1,5 @@
/*
* Copyright (c) 2020-2021, NVIDIA CORPORATION.
* Copyright (c) 2020-2022, NVIDIA CORPORATION.
*
* Licensed under the Apache License, Version 2.0 (the "License");
* you may not use this file except in compliance with the License.
Expand Down Expand Up @@ -79,7 +79,7 @@ data_type arrow_to_cudf_type(arrow::DataType const& arrow_type)
case arrow::Type::LIST: return data_type(type_id::LIST);
case arrow::Type::DECIMAL: {
auto const type = static_cast<arrow::Decimal128Type const*>(&arrow_type);
return data_type{type_id::DECIMAL64, -type->scale()};
return data_type{type_id::DECIMAL128, -type->scale()};
}
case arrow::Type::STRUCT: return data_type(type_id::STRUCT);
default: CUDF_FAIL("Unsupported type_id conversion to cudf");
Expand Down Expand Up @@ -177,35 +177,27 @@ std::unique_ptr<column> get_column(arrow::Array const& array,
rmm::mr::device_memory_resource* mr);

template <>
std::unique_ptr<column> dispatch_to_cudf_column::operator()<numeric::decimal64>(
std::unique_ptr<column> dispatch_to_cudf_column::operator()<numeric::decimal128>(
arrow::Array const& array,
data_type type,
bool skip_mask,
rmm::cuda_stream_view stream,
rmm::mr::device_memory_resource* mr)
{
using DeviceType = int64_t;
using DeviceType = __int128_t;

auto constexpr BIT_WIDTH_RATIO = 2; // Array::Type:type::DECIMAL (128) / int64_t
auto data_buffer = array.data()->buffers[1];
auto const num_rows = static_cast<size_type>(array.length());

rmm::device_uvector<DeviceType> buf(num_rows * BIT_WIDTH_RATIO, stream);
rmm::device_uvector<DeviceType> out_buf(num_rows, stream, mr);
auto data_buffer = array.data()->buffers[1];
auto const num_rows = static_cast<size_type>(array.length());
auto col = make_fixed_width_column(type, num_rows, mask_state::UNALLOCATED, stream, mr);
auto mutable_column_view = col->mutable_view();

CUDA_TRY(cudaMemcpyAsync(
reinterpret_cast<uint8_t*>(buf.data()),
mutable_column_view.data<DeviceType>(),
reinterpret_cast<const uint8_t*>(data_buffer->address()) + array.offset() * sizeof(DeviceType),
buf.size() * sizeof(DeviceType),
sizeof(DeviceType) * num_rows,
cudaMemcpyDefault,
stream.value()));

auto every_other = [] __device__(size_type i) { return 2 * i; };
auto gather_map = cudf::detail::make_counting_transform_iterator(0, every_other);

thrust::gather(
rmm::exec_policy(stream), gather_map, gather_map + num_rows, buf.data(), out_buf.data());

auto null_mask = [&] {
if (not skip_mask and array.null_bitmap_data()) {
auto temp_mask = get_mask_buffer(array, stream, mr);
Expand All @@ -221,7 +213,8 @@ std::unique_ptr<column> dispatch_to_cudf_column::operator()<numeric::decimal64>(
return rmm::device_buffer{};
}();

return std::make_unique<cudf::column>(type, num_rows, out_buf.release(), std::move(null_mask));
col->set_null_mask(std::move(null_mask));
return col;
}

template <>
Expand Down
36 changes: 35 additions & 1 deletion cpp/src/interop/to_arrow.cu
View file
Original file line number Diff line number Diff line change
@@ -1,5 +1,5 @@
/*
* Copyright (c) 2020-2021, NVIDIA CORPORATION.
* Copyright (c) 2020-2022, NVIDIA CORPORATION.
*
* Licensed under the Apache License, Version 2.0 (the "License");
* you may not use this file except in compliance with the License.
Expand Down Expand Up @@ -176,6 +176,40 @@ std::shared_ptr<arrow::Array> dispatch_to_arrow::operator()<numeric::decimal64>(
return std::make_shared<arrow::Decimal128Array>(data);
}

template <>
std::shared_ptr<arrow::Array> dispatch_to_arrow::operator()<numeric::decimal128>(
column_view input,
cudf::type_id,
column_metadata const&,
arrow::MemoryPool* ar_mr,
rmm::cuda_stream_view stream)
{
using DeviceType = __int128_t;

rmm::device_uvector<DeviceType> buf(input.size(), stream);

thrust::copy(rmm::exec_policy(stream), //
input.begin<DeviceType>(),
input.end<DeviceType>(),
buf.begin());

auto const buf_size_in_bytes = buf.size() * sizeof(DeviceType);
auto data_buffer = allocate_arrow_buffer(buf_size_in_bytes, ar_mr);

CUDA_TRY(cudaMemcpyAsync(data_buffer->mutable_data(),
buf.data(),
buf_size_in_bytes,
cudaMemcpyDeviceToHost,
stream.value()));

auto type = arrow::decimal(18, -input.type().scale());
auto mask = fetch_mask_buffer(input, ar_mr, stream);
auto buffers = std::vector<std::shared_ptr<arrow::Buffer>>{mask, std::move(data_buffer)};
auto data = std::make_shared<arrow::ArrayData>(type, input.size(), buffers);

return std::make_shared<arrow::Decimal128Array>(data);
}

template <>
std::shared_ptr<arrow::Array> dispatch_to_arrow::operator()<bool>(column_view input,
cudf::type_id id,
Expand Down
26 changes: 25 additions & 1 deletion cpp/tests/interop/arrow_utils.hpp
View file
Original file line number Diff line number Diff line change
@@ -1,5 +1,5 @@
/*
* Copyright (c) 2020, NVIDIA CORPORATION.
* Copyright (c) 2020-2022, NVIDIA CORPORATION.
*
* Licensed under the Apache License, Version 2.0 (the "License");
* you may not use this file except in compliance with the License.
Expand Down Expand Up @@ -176,3 +176,27 @@ std::shared_ptr<arrow::Array> get_arrow_list_array(

std::pair<std::unique_ptr<cudf::table>, std::shared_ptr<arrow::Table>> get_tables(
cudf::size_type length = 10000);

template <typename T>
[[nodiscard]] auto make_decimal128_arrow_array(std::vector<T> const& data,
std::optional<std::vector<int>> const& validity,
int32_t scale) -> std::shared_ptr<arrow::Array>
{
auto constexpr BIT_WIDTH_RATIO = sizeof(__int128_t) / sizeof(T);

std::shared_ptr<arrow::Array> arr;
arrow::Decimal128Builder decimal_builder(arrow::decimal(18, -scale),
arrow::default_memory_pool());

for (T i = 0; i < static_cast<T>(data.size() / BIT_WIDTH_RATIO); ++i) {
if (validity.has_value() and not validity.value()[i]) {
decimal_builder.AppendNull();
} else {
decimal_builder.Append(reinterpret_cast<const uint8_t*>(data.data() + BIT_WIDTH_RATIO * i));
}
}

CUDF_EXPECTS(decimal_builder.Finish(&arr).ok(), "Failed to build array");

return arr;
}
81 changes: 28 additions & 53 deletions cpp/tests/interop/from_arrow_test.cpp
View file
Original file line number Diff line number Diff line change
@@ -1,5 +1,5 @@
/*
* Copyright (c) 2020-2021, NVIDIA CORPORATION.
* Copyright (c) 2020-2022, NVIDIA CORPORATION.
*
* Licensed under the Apache License, Version 2.0 (the "License");
* you may not use this file except in compliance with the License.
Expand Down Expand Up @@ -354,21 +354,16 @@ TEST_P(FromArrowTestSlice, SliceTest)
template <typename T>
using fp_wrapper = cudf::test::fixed_point_column_wrapper<T>;

TEST_F(FromArrowTest, FixedPointTable)
TEST_F(FromArrowTest, FixedPoint128Table)
{
using namespace numeric;
auto constexpr BIT_WIDTH_RATIO = 2; // Array::Type:type::DECIMAL (128) / int64_t

for (auto const i : {3, 2, 1, 0, -1, -2, -3}) {
auto const data = std::vector<int64_t>{1, 0, 2, 0, 3, 0, 4, 0, 5, 0, 6, 0};
auto const col = fp_wrapper<int64_t>({1, 2, 3, 4, 5, 6}, scale_type{i});
for (auto const scale : {3, 2, 1, 0, -1, -2, -3}) {
auto const data = std::vector<__int128_t>{1, 2, 3, 4, 5, 6};
auto const col = fp_wrapper<__int128_t>(data.cbegin(), data.cend(), scale_type{scale});
auto const expected = cudf::table_view({col});

std::shared_ptr<arrow::Array> arr;
arrow::Decimal128Builder decimal_builder(arrow::decimal(10, -i), arrow::default_memory_pool());
decimal_builder.AppendValues(reinterpret_cast<const uint8_t*>(data.data()),
data.size() / BIT_WIDTH_RATIO);
CUDF_EXPECTS(decimal_builder.Finish(&arr).ok(), "Failed to build array");
auto const arr = make_decimal128_arrow_array(data, std::nullopt, scale);

auto const field = arrow::field("a", arr->type());
auto const schema_vector = std::vector<std::shared_ptr<arrow::Field>>({field});
Expand All @@ -381,24 +376,18 @@ TEST_F(FromArrowTest, FixedPointTable)
}
}

TEST_F(FromArrowTest, FixedPointTableLarge)
TEST_F(FromArrowTest, FixedPoint128TableLarge)
{
using namespace numeric;
auto constexpr BIT_WIDTH_RATIO = 2; // Array::Type:type::DECIMAL (128) / int64_t
auto constexpr NUM_ELEMENTS = 1000;
auto constexpr NUM_ELEMENTS = 1000;

for (auto const i : {3, 2, 1, 0, -1, -2, -3}) {
auto every_other = [](auto i) { return i % BIT_WIDTH_RATIO ? 0 : i / BIT_WIDTH_RATIO; };
auto transform = cudf::detail::make_counting_transform_iterator(BIT_WIDTH_RATIO, every_other);
auto const data = std::vector<int64_t>(transform, transform + NUM_ELEMENTS * BIT_WIDTH_RATIO);
auto iota = thrust::make_counting_iterator(1);
auto const col = fp_wrapper<int64_t>(iota, iota + NUM_ELEMENTS, scale_type{i});
for (auto const scale : {3, 2, 1, 0, -1, -2, -3}) {
auto iota = thrust::make_counting_iterator(1);
auto const data = std::vector<__int128_t>(iota, iota + NUM_ELEMENTS);
auto const col = fp_wrapper<__int128_t>(iota, iota + NUM_ELEMENTS, scale_type{scale});
auto const expected = cudf::table_view({col});

std::shared_ptr<arrow::Array> arr;
arrow::Decimal128Builder decimal_builder(arrow::decimal(10, -i), arrow::default_memory_pool());
decimal_builder.AppendValues(reinterpret_cast<const uint8_t*>(data.data()), NUM_ELEMENTS);
CUDF_EXPECTS(decimal_builder.Finish(&arr).ok(), "Failed to build array");
auto const arr = make_decimal128_arrow_array(data, std::nullopt, scale);

auto const field = arrow::field("a", arr->type());
auto const schema_vector = std::vector<std::shared_ptr<arrow::Field>>({field});
Expand All @@ -411,25 +400,18 @@ TEST_F(FromArrowTest, FixedPointTableLarge)
}
}

TEST_F(FromArrowTest, FixedPointTableNulls)
TEST_F(FromArrowTest, FixedPoint128TableNulls)
{
using namespace numeric;
auto constexpr BIT_WIDTH_RATIO = 2; // Array::Type:type::DECIMAL (128) / int64_t

for (auto const i : {3, 2, 1, 0, -1, -2, -3}) {
auto const data = std::vector<int64_t>{1, 0, 2, 0, 3, 0, 4, 0, 5, 0, 6, 0};
for (auto const scale : {3, 2, 1, 0, -1, -2, -3}) {
auto const data = std::vector<__int128_t>{1, 2, 3, 4, 5, 6, 0, 0};
auto const validity = std::vector<int32_t>{1, 1, 1, 1, 1, 1, 0, 0};
auto const col =
fp_wrapper<int64_t>({1, 2, 3, 4, 5, 6, 0, 0}, {1, 1, 1, 1, 1, 1, 0, 0}, scale_type{i});
fp_wrapper<__int128_t>({1, 2, 3, 4, 5, 6, 0, 0}, {1, 1, 1, 1, 1, 1, 0, 0}, scale_type{scale});
auto const expected = cudf::table_view({col});

std::shared_ptr<arrow::Array> arr;
arrow::Decimal128Builder decimal_builder(arrow::decimal(10, -i), arrow::default_memory_pool());
decimal_builder.AppendValues(reinterpret_cast<const uint8_t*>(data.data()),
data.size() / BIT_WIDTH_RATIO);
decimal_builder.AppendNull();
decimal_builder.AppendNull();

CUDF_EXPECTS(decimal_builder.Finish(&arr).ok(), "Failed to build array");
auto const arr = make_decimal128_arrow_array(data, validity, scale);

auto const field = arrow::field("a", arr->type());
auto const schema_vector = std::vector<std::shared_ptr<arrow::Field>>({field});
Expand All @@ -442,28 +424,21 @@ TEST_F(FromArrowTest, FixedPointTableNulls)
}
}

TEST_F(FromArrowTest, FixedPointTableNullsLarge)
TEST_F(FromArrowTest, FixedPoint128TableNullsLarge)
{
using namespace numeric;
auto constexpr BIT_WIDTH_RATIO = 2; // Array::Type:type::DECIMAL (128) / int64_t
auto constexpr NUM_ELEMENTS = 1000;
auto constexpr NUM_ELEMENTS = 1000;

for (auto const i : {3, 2, 1, 0, -1, -2, -3}) {
auto every_other = [](auto i) { return i % BIT_WIDTH_RATIO ? 0 : i / BIT_WIDTH_RATIO; };
auto transform = cudf::detail::make_counting_transform_iterator(BIT_WIDTH_RATIO, every_other);
auto const data = std::vector<int64_t>(transform, transform + NUM_ELEMENTS * BIT_WIDTH_RATIO);
for (auto const scale : {3, 2, 1, 0, -1, -2, -3}) {
auto every_other = [](auto i) { return i % 2 ? 0 : 1; };
auto validity = cudf::detail::make_counting_transform_iterator(0, every_other);
auto iota = thrust::make_counting_iterator(1);
auto const col = fp_wrapper<int64_t>(iota, iota + NUM_ELEMENTS, transform, scale_type{i});
auto const data = std::vector<__int128_t>(iota, iota + NUM_ELEMENTS);
auto const col = fp_wrapper<__int128_t>(iota, iota + NUM_ELEMENTS, validity, scale_type{scale});
auto const expected = cudf::table_view({col});

std::shared_ptr<arrow::Array> arr;
arrow::Decimal128Builder decimal_builder(arrow::decimal(10, -i), arrow::default_memory_pool());
for (int64_t i = 0; i < NUM_ELEMENTS / BIT_WIDTH_RATIO; ++i) {
decimal_builder.Append(reinterpret_cast<const uint8_t*>(data.data() + 4 * i));
decimal_builder.AppendNull();
}

CUDF_EXPECTS(decimal_builder.Finish(&arr).ok(), "Failed to build array");
auto const arr = make_decimal128_arrow_array(
data, std::vector<int32_t>(validity, validity + NUM_ELEMENTS), scale);

auto const field = arrow::field("a", arr->type());
auto const schema_vector = std::vector<std::shared_ptr<arrow::Field>>({field});
Expand Down
Loading

0 comments on commit 45c20d1

Please sign in to comment.