Skip to content

Commit

Permalink
Added decimal convertions from/to python
Browse files Browse the repository at this point in the history 
Signed-off-by: Dmitry Chigarev <dmitry.chigarev@intel.com>
  • Loading branch information
@dchigarev
dchigarev committed Dec 28, 2020
1 parent e41670f commit a4c8797
Show file tree
Hide file tree
Showing 5 changed files with 145 additions and 39 deletions.
38 changes: 10 additions & 28 deletions cpp/src/arrow/python/arrow_to_pandas.cc
View file
Original file line number Diff line number Diff line change
Expand Up @@ -167,6 +167,9 @@ static inline bool ListTypeSupported(const DataType& type) {
case Type::UINT64:
case Type::FLOAT:
case Type::DOUBLE:
case Type::DECIMAL16:
case Type::DECIMAL32:
case Type::DECIMAL64:
case Type::DECIMAL128:
case Type::DECIMAL256:
case Type::BINARY:
Expand Down Expand Up @@ -1021,7 +1024,7 @@ struct ObjectWriterVisitor {
}

template <typename Type>
enable_if_t<is_base_binary_type<Type>::value || is_fixed_size_binary_type<Type>::value,
enable_if_t<(is_base_binary_type<Type>::value || is_fixed_size_binary_type<Type>::value) && !is_decimal_type<Type>::value,
Status>
Visit(const Type& type) {
auto WrapValue = [](const util::string_view& view, PyObject** out) {
Expand Down Expand Up @@ -1094,40 +1097,16 @@ struct ObjectWriterVisitor {
return Status::OK();
}

Status Visit(const Decimal128Type& type) {
template <uint32_t width>
Status Visit(const BaseDecimalType<width>& type) {
OwnedRef decimal;
OwnedRef Decimal;
RETURN_NOT_OK(internal::ImportModule("decimal", &decimal));
RETURN_NOT_OK(internal::ImportFromModule(decimal.obj(), "Decimal", &Decimal));
PyObject* decimal_constructor = Decimal.obj();

for (int c = 0; c < data.num_chunks(); c++) {
const auto& arr = checked_cast<const arrow::Decimal128Array&>(*data.chunk(c));

for (int64_t i = 0; i < arr.length(); ++i) {
if (arr.IsNull(i)) {
Py_INCREF(Py_None);
*out_values++ = Py_None;
} else {
*out_values++ =
internal::DecimalFromString(decimal_constructor, arr.FormatValue(i));
RETURN_IF_PYERROR();
}
}
}

return Status::OK();
}

Status Visit(const Decimal256Type& type) {
OwnedRef decimal;
OwnedRef Decimal;
RETURN_NOT_OK(internal::ImportModule("decimal", &decimal));
RETURN_NOT_OK(internal::ImportFromModule(decimal.obj(), "Decimal", &Decimal));
PyObject* decimal_constructor = Decimal.obj();

for (int c = 0; c < data.num_chunks(); c++) {
const auto& arr = checked_cast<const arrow::Decimal256Array&>(*data.chunk(c));
const auto& arr = checked_cast<const arrow::BaseDecimalArray<width>&>(*data.chunk(c));

for (int64_t i = 0; i < arr.length(); ++i) {
if (arr.IsNull(i)) {
Expand Down Expand Up @@ -1871,6 +1850,9 @@ static Status GetPandasWriterType(const ChunkedArray& data, const PandasOptions&
case Type::STRUCT: // fall through
case Type::TIME32: // fall through
case Type::TIME64: // fall through
case Type::DECIMAL16: // fall through
case Type::DECIMAL32: // fall through
case Type::DECIMAL64: // fall through
case Type::DECIMAL128: // fall through
case Type::DECIMAL256: // fall through
*output_type = PandasWriter::OBJECT;
Expand Down
30 changes: 30 additions & 0 deletions cpp/src/arrow/python/decimal.cc
View file
Original file line number Diff line number Diff line change
Expand Up @@ -166,6 +166,36 @@ Status InternalDecimalFromPyObject(PyObject* obj, const DecimalType& arrow_type,

} // namespace

Status DecimalFromPythonDecimal(PyObject* python_decimal, const DecimalType& arrow_type,
Decimal16* out) {
return InternalDecimalFromPythonDecimal(python_decimal, arrow_type, out);
}

Status DecimalFromPyObject(PyObject* obj, const DecimalType& arrow_type,
Decimal16* out) {
return InternalDecimalFromPyObject(obj, arrow_type, out);
}

Status DecimalFromPythonDecimal(PyObject* python_decimal, const DecimalType& arrow_type,
Decimal32* out) {
return InternalDecimalFromPythonDecimal(python_decimal, arrow_type, out);
}

Status DecimalFromPyObject(PyObject* obj, const DecimalType& arrow_type,
Decimal32* out) {
return InternalDecimalFromPyObject(obj, arrow_type, out);
}

Status DecimalFromPythonDecimal(PyObject* python_decimal, const DecimalType& arrow_type,
Decimal64* out) {
return InternalDecimalFromPythonDecimal(python_decimal, arrow_type, out);
}

Status DecimalFromPyObject(PyObject* obj, const DecimalType& arrow_type,
Decimal64* out) {
return InternalDecimalFromPyObject(obj, arrow_type, out);
}

Status DecimalFromPythonDecimal(PyObject* python_decimal, const DecimalType& arrow_type,
Decimal128* out) {
return InternalDecimalFromPythonDecimal(python_decimal, arrow_type, out);
Expand Down
57 changes: 57 additions & 0 deletions cpp/src/arrow/python/decimal.h
View file
Original file line number Diff line number Diff line change
Expand Up @@ -24,6 +24,12 @@

namespace arrow {

template<uint32_t width>
class DecimalAnyWidth;

using Decimal16 = DecimalAnyWidth<16>;
using Decimal32 = DecimalAnyWidth<32>;
using Decimal64 = DecimalAnyWidth<64>;
class Decimal128;
class Decimal256;

Expand Down Expand Up @@ -56,6 +62,57 @@ ARROW_PYTHON_EXPORT
PyObject* DecimalFromString(PyObject* decimal_constructor,
const std::string& decimal_string);

// \brief Convert a Python decimal to an Arrow Decimal16 object
// \param[in] python_decimal A Python decimal.Decimal instance
// \param[in] arrow_type An instance of arrow::DecimalType
// \param[out] out A pointer to a Decimal16
// \return The status of the operation
ARROW_PYTHON_EXPORT
Status DecimalFromPythonDecimal(PyObject* python_decimal, const DecimalType& arrow_type,
Decimal16* out);

// \brief Convert a Python object to an Arrow Decimal16 object
// \param[in] python_decimal A Python int or decimal.Decimal instance
// \param[in] arrow_type An instance of arrow::DecimalType
// \param[out] out A pointer to a Decimal16
// \return The status of the operation
ARROW_PYTHON_EXPORT
Status DecimalFromPyObject(PyObject* obj, const DecimalType& arrow_type, Decimal16* out);

// \brief Convert a Python decimal to an Arrow Decimal32 object
// \param[in] python_decimal A Python decimal.Decimal instance
// \param[in] arrow_type An instance of arrow::DecimalType
// \param[out] out A pointer to a Decimal32
// \return The status of the operation
ARROW_PYTHON_EXPORT
Status DecimalFromPythonDecimal(PyObject* python_decimal, const DecimalType& arrow_type,
Decimal32* out);

// \brief Convert a Python object to an Arrow Decimal32 object
// \param[in] python_decimal A Python int or decimal.Decimal instance
// \param[in] arrow_type An instance of arrow::DecimalType
// \param[out] out A pointer to a Decimal32
// \return The status of the operation
ARROW_PYTHON_EXPORT
Status DecimalFromPyObject(PyObject* obj, const DecimalType& arrow_type, Decimal32* out);

// \brief Convert a Python decimal to an Arrow Decimal64 object
// \param[in] python_decimal A Python decimal.Decimal instance
// \param[in] arrow_type An instance of arrow::DecimalType
// \param[out] out A pointer to a Decimal64
// \return The status of the operation
ARROW_PYTHON_EXPORT
Status DecimalFromPythonDecimal(PyObject* python_decimal, const DecimalType& arrow_type,
Decimal64* out);

// \brief Convert a Python object to an Arrow Decimal64 object
// \param[in] python_decimal A Python int or decimal.Decimal instance
// \param[in] arrow_type An instance of arrow::DecimalType
// \param[out] out A pointer to a Decimal64
// \return The status of the operation
ARROW_PYTHON_EXPORT
Status DecimalFromPyObject(PyObject* obj, const DecimalType& arrow_type, Decimal64* out);

// \brief Convert a Python decimal to an Arrow Decimal128 object
// \param[in] python_decimal A Python decimal.Decimal instance
// \param[in] arrow_type An instance of arrow::DecimalType
Expand Down
41 changes: 30 additions & 11 deletions cpp/src/arrow/python/python_test.cc
View file
Original file line number Diff line number Diff line change
Expand Up @@ -28,6 +28,7 @@
#include "arrow/table.h"
#include "arrow/testing/gtest_util.h"
#include "arrow/util/decimal.h"
#include "arrow/util/decimal_type_traits.h"
#include "arrow/util/optional.h"

#include "arrow/python/arrow_to_pandas.h"
Expand Down Expand Up @@ -358,6 +359,12 @@ void DecimalTestFromPythonDecimalRescale(std::shared_ptr<DataType> type,
TEST_F(DecimalTest, FromPythonDecimalRescaleNotTruncateable) {
// We fail when truncating values that would lose data if cast to a decimal type with
// lower scale
DecimalTestFromPythonDecimalRescale<Decimal16>(::arrow::decimal16(5, 2),
this->CreatePythonDecimal("1.001"), {});
DecimalTestFromPythonDecimalRescale<Decimal32>(::arrow::decimal32(10, 2),
this->CreatePythonDecimal("1.001"), {});
DecimalTestFromPythonDecimalRescale<Decimal64>(::arrow::decimal64(10, 2),
this->CreatePythonDecimal("1.001"), {});
DecimalTestFromPythonDecimalRescale<Decimal128>(::arrow::decimal128(10, 2),
this->CreatePythonDecimal("1.001"), {});
DecimalTestFromPythonDecimalRescale<Decimal256>(::arrow::decimal256(10, 2),
Expand All @@ -367,32 +374,44 @@ TEST_F(DecimalTest, FromPythonDecimalRescaleNotTruncateable) {
TEST_F(DecimalTest, FromPythonDecimalRescaleTruncateable) {
// We allow truncation of values that do not lose precision when dividing by 10 * the
// difference between the scales, e.g., 1.000 -> 1.00
DecimalTestFromPythonDecimalRescale<Decimal16>(
::arrow::decimal16(5, 2), this->CreatePythonDecimal("1.000"), 100);
DecimalTestFromPythonDecimalRescale<Decimal32>(
::arrow::decimal32(10, 2), this->CreatePythonDecimal("1.000"), 100);
DecimalTestFromPythonDecimalRescale<Decimal64>(
::arrow::decimal64(10, 2), this->CreatePythonDecimal("1.000"), 100);
DecimalTestFromPythonDecimalRescale<Decimal128>(
::arrow::decimal128(10, 2), this->CreatePythonDecimal("1.000"), 100);
DecimalTestFromPythonDecimalRescale<Decimal256>(
::arrow::decimal256(10, 2), this->CreatePythonDecimal("1.000"), 100);
}

TEST_F(DecimalTest, FromPythonNegativeDecimalRescale) {
DecimalTestFromPythonDecimalRescale<Decimal128>(
::arrow::decimal16(5, 4), this->CreatePythonDecimal("-1.000"), -10000);
DecimalTestFromPythonDecimalRescale<Decimal32>(
::arrow::decimal32(10, 9), this->CreatePythonDecimal("-1.000"), -1000000000);
DecimalTestFromPythonDecimalRescale<Decimal64>(
::arrow::decimal64(10, 9), this->CreatePythonDecimal("-1.000"), -1000000000);
DecimalTestFromPythonDecimalRescale<Decimal128>(
::arrow::decimal128(10, 9), this->CreatePythonDecimal("-1.000"), -1000000000);
DecimalTestFromPythonDecimalRescale<Decimal256>(
::arrow::decimal256(10, 9), this->CreatePythonDecimal("-1.000"), -1000000000);
}

TEST_F(DecimalTest, Decimal128FromPythonInteger) {
Decimal128 value;
OwnedRef python_long(PyLong_FromLong(42));
auto type = ::arrow::decimal128(10, 2);
const auto& decimal_type = checked_cast<const DecimalType&>(*type);
ASSERT_OK(internal::DecimalFromPyObject(python_long.obj(), decimal_type, &value));
ASSERT_EQ(4200, value);
}
template<typename T>
class DecimalTestConversion : public testing::Test {};
using DecimalTypes = ::testing::Types<DecimalTypeTraits<16>, DecimalTypeTraits<32>, DecimalTypeTraits<64>, DecimalTypeTraits<128>, DecimalTypeTraits<256>>;

TEST_F(DecimalTest, Decimal256FromPythonInteger) {
Decimal256 value;
TYPED_TEST_SUITE(DecimalTestConversion, DecimalTypes);

TYPED_TEST(DecimalTestConversion, Basics) {
using TypeClass = typename TypeParam::TypeClass;
using ValueType = typename TypeParam::ValueType;

ValueType value;
OwnedRef python_long(PyLong_FromLong(42));
auto type = ::arrow::decimal256(10, 2);
auto type = std::make_shared<TypeClass>(5, 2);
const auto& decimal_type = checked_cast<const DecimalType&>(*type);
ASSERT_OK(internal::DecimalFromPyObject(python_long.obj(), decimal_type, &value));
ASSERT_EQ(4200, value);
Expand Down
18 changes: 18 additions & 0 deletions cpp/src/arrow/python/python_to_arrow.cc
View file
Original file line number Diff line number Diff line change
Expand Up @@ -164,6 +164,24 @@ class PyValue {
return value;
}

static Result<Decimal16> Convert(const Decimal16Type* type, const O&, I obj) {
Decimal16 value;
RETURN_NOT_OK(internal::DecimalFromPyObject(obj, *type, &value));
return value;
}

static Result<Decimal32> Convert(const Decimal32Type* type, const O&, I obj) {
Decimal32 value;
RETURN_NOT_OK(internal::DecimalFromPyObject(obj, *type, &value));
return value;
}

static Result<Decimal64> Convert(const Decimal64Type* type, const O&, I obj) {
Decimal64 value;
RETURN_NOT_OK(internal::DecimalFromPyObject(obj, *type, &value));
return value;
}

static Result<Decimal128> Convert(const Decimal128Type* type, const O&, I obj) {
Decimal128 value;
RETURN_NOT_OK(internal::DecimalFromPyObject(obj, *type, &value));
Expand Down

0 comments on commit a4c8797

Please sign in to comment.