This change enables support for DECIMALs backed by BYTE_ARRAYs on disk. It does this by creating a TransferFunctor routine that transforms a ByteArrayType to an ::arrow::Decimal128Type.

The routine does this by:
1. Creating an arrow::BinaryArray from the RecordReader's builder
2. Allocating a buffer for the arrow::Decimal128Array
3. Converting the big-endian bytes in each BinaryArray entry to two integers
   representing the high and low bits of each decimal value.
This commit is contained in:
Ted Haining authored and Wes McKinney committed 2018-09-28 05:52:46 -04:00
1 parent c3c024ecf2
commit 30f3a27801
2 files changed
+68 -5

No files matched your search

@@ -2316,11 +2316,11 @@ TEST(TestArrowReaderAdHoc, Int96BadMemoryAccess) {
ASSERT_OK_NO_THROW(arrow_reader->ReadTable(&table));
}
class TestArrowReaderAdHocSpark
class TestArrowReaderAdHocSparkAndHvr
: public ::testing::TestWithParam<
std::tuple<std::string, std::shared_ptr<::DataType>>> {};
TEST_P(TestArrowReaderAdHocSpark, ReadDecimals) {
TEST_P(TestArrowReaderAdHocSparkAndHvr, ReadDecimals) {
std::string path(test::get_data_dir());
std::string filename;
@@ -2364,12 +2364,13 @@ TEST_P(TestArrowReaderAdHocSpark, ReadDecimals) {
}
INSTANTIATE_TEST_CASE_P(
ReadDecimals, TestArrowReaderAdHocSpark,
ReadDecimals, TestArrowReaderAdHocSparkAndHvr,
::testing::Values(
std::make_tuple("int32_decimal.parquet", ::arrow::decimal(4, 2)),
std::make_tuple("int64_decimal.parquet", ::arrow::decimal(10, 2)),
std::make_tuple("fixed_length_decimal.parquet", ::arrow::decimal(25, 2)),
std::make_tuple("fixed_length_decimal_legacy.parquet", ::arrow::decimal(13, 2))));
std::make_tuple("fixed_length_decimal_legacy.parquet", ::arrow::decimal(13, 2)),
std::make_tuple("byte_array_decimal.parquet", ::arrow::decimal(4, 2))));
} // namespace arrow
+63 -1
View File
@@ -1221,6 +1221,64 @@ struct TransferFunctor<::arrow::Decimal128Type, FLBAType> {
}
};
/// \brief Convert an arrow::BinaryArray to an arrow::Decimal128Array
/// We do this by:
/// 1. Creating an arrow::BinaryArray from the RecordReader's builder
/// 2. Allocating a buffer for the arrow::Decimal128Array
/// 3. Converting the big-endian bytes in each BinaryArray entry to two integers
/// representing the high and low bits of each decimal value.
template <>
struct TransferFunctor<::arrow::Decimal128Type, ByteArrayType> {
Status operator()(RecordReader* reader, MemoryPool* pool,
const std::shared_ptr<::arrow::DataType>& type,
std::shared_ptr<Array>* out) {
DCHECK_EQ(type->id(), ::arrow::Type::DECIMAL);
// Finish the built data into a temporary array
std::shared_ptr<Array> array;
RETURN_NOT_OK(reader->builder()->Finish(&array));
const auto& binary_array = static_cast<const ::arrow::BinaryArray&>(*array);
const int64_t length = binary_array.length();
const auto& decimal_type = static_cast<const ::arrow::Decimal128Type&>(*type);
const int64_t type_length = decimal_type.byte_width();
std::shared_ptr<Buffer> data;
RETURN_NOT_OK(::arrow::AllocateBuffer(pool, length * type_length, &data));
// raw bytes that we can write to
uint8_t* out_ptr = data->mutable_data();
const int64_t null_count = binary_array.null_count();
// convert each BinaryArray value to valid decimal bytes
for (int64_t i = 0; i < length; i++, out_ptr += type_length) {
int32_t record_len = 0;
const uint8_t* record_loc = binary_array.GetValue(i, &record_len);
if ((record_len < 0) || (record_len > type_length)) {
return Status::Invalid("Invalid BYTE_ARRAY size");
}
auto out_ptr_view = reinterpret_cast<uint64_t*>(out_ptr);
out_ptr_view[0] = 0;
out_ptr_view[1] = 0;
// only convert rows that are not null if there are nulls, or
// all rows, if there are not
if (((null_count > 0) && !binary_array.IsNull(i)) || (null_count <= 0)) {
RawBytesToDecimalBytes(record_loc, record_len, out_ptr);
}
}
*out = std::make_shared<::arrow::Decimal128Array>(
type, length, data, binary_array.null_bitmap(), null_count);
return Status::OK();
}
};
/// \brief Convert an Int32 or Int64 array into a Decimal128Array
/// The parquet spec allows systems to write decimals in int32, int64 if the values are
/// small enough to fit in less 4 bytes or less than 8 bytes, respectively.
@@ -1353,12 +1411,16 @@ Status PrimitiveImpl::NextBatch(int64_t records_to_read, std::shared_ptr<Array>*
case ::parquet::Type::INT64: {
TRANSFER_DATA(::arrow::Decimal128Type, Int64Type);
} break;
case ::parquet::Type::BYTE_ARRAY: {
TRANSFER_DATA(::arrow::Decimal128Type, ByteArrayType);
} break;
case ::parquet::Type::FIXED_LEN_BYTE_ARRAY: {
TRANSFER_DATA(::arrow::Decimal128Type, FLBAType);
} break;
default:
return Status::Invalid(
"Physical type for decimal must be int32, int64, or fixed length binary");
"Physical type for decimal must be int32, int64, byte array, or fixed "
"length binary");
}
} break;
case ::arrow::Type::TIMESTAMP: {