arrow_utils.h

Go to the documentation of this file.

 
#ifndef UTILS_ARROW_UTILS_H_
#define UTILS_ARROW_UTILS_H_
 
#include <arrow/api.h>
#include <arrow/util/value_parsing.h>
#include <memory>
#include "flex/utils/property/types.h"
#include "glog/logging.h"
 
namespace gs {
 
// arrow related;
 
class LDBCTimeStampParser : public arrow::TimestampParser {
 public:
 LDBCTimeStampParser() = default;
 
 ~LDBCTimeStampParser() override {}
 
 bool operator()(const char* s, size_t length, arrow::TimeUnit::type out_unit,
 int64_t* out,
 bool* out_zone_offset_present = NULLPTR) const override {
 using seconds_type = std::chrono::duration<arrow::TimestampType::c_type>;
 
 // We allow the following zone offset formats:
 // - (none)
 // - Z
 // - [+-]HH(:?MM)?
 //
 // We allow the following formats for all units:
 // - "YYYY-MM-DD"
 // - "YYYY-MM-DD[ T]hhZ?"
 // - "YYYY-MM-DD[ T]hh:mmZ?"
 // - "YYYY-MM-DD[ T]hh:mm:ssZ?"
 //
 // We allow the following formats for unit == MILLI, MICRO, or NANO:
 // - "YYYY-MM-DD[ T]hh:mm:ss.s{1,3}Z?"
 //
 // We allow the following formats for unit == MICRO, or NANO:
 // - "YYYY-MM-DD[ T]hh:mm:ss.s{4,6}Z?"
 //
 // We allow the following formats for unit == NANO:
 // - "YYYY-MM-DD[ T]hh:mm:ss.s{7,9}Z?"
 //
 // UTC is always assumed, and the DataType's timezone is ignored.
 //
 
 if (ARROW_PREDICT_FALSE(length < 10))
 return false;
 
 seconds_type seconds_since_epoch;
#if defined(ARROW_VERSION) && ARROW_VERSION < 15000000
 if (ARROW_PREDICT_FALSE(!arrow::internal::detail::ParseYYYY_MM_DD(
 s, &seconds_since_epoch))) {
#else
 if (ARROW_PREDICT_FALSE(
 !arrow::internal::ParseYYYY_MM_DD(s, &seconds_since_epoch))) {
#endif
 return false;
 }
 
 if (length == 10) {
 *out =
 arrow::util::CastSecondsToUnit(out_unit, seconds_since_epoch.count());
 return true;
 }
 
 if (ARROW_PREDICT_FALSE(s[10] != ' ') &&
 ARROW_PREDICT_FALSE(s[10] != 'T')) {
 return false;
 }
 
 // In the implementation of arrow ISO8601 timestamp parser, the zone offset
 // is set to true if the input string contains a zone offset. However, we
 // parse the zone offset here but don't set the boolean flag.
 // https://github.com/apache/arrow/blob/3e7ae5340a123c1040f98f1c36687b81362fab52/cpp/src/arrow/csv/converter.cc#L373
 // The reason is that, if we want the zone offset to be set, we need to
 // to declare the zone offset in the schema and construct TimeStampType with
 // that offset. However, we just want to parse the timestamp string and
 // convert it to a timestamp value, we have no assumption of the local time
 // zone, and we don't require the zone offset to be set in the schema.
 // Same for following commented code.
 //-------------------------------------------------------------------------
 // if (out_zone_offset_present) {
 // *out_zone_offset_present = false;
 // }
 //-------------------------------------------------------------------------
 
 seconds_type zone_offset(0);
 if (s[length - 1] == 'Z') {
 --length;
 // if (out_zone_offset_present)
 // *out_zone_offset_present = true;
 } else if (s[length - 3] == '+' || s[length - 3] == '-') {
 // [+-]HH
 length -= 3;
 if (ARROW_PREDICT_FALSE(!arrow::internal::detail::ParseHH(
 s + length + 1, &zone_offset))) {
 return false;
 }
 if (s[length] == '+')
 zone_offset *= -1;
 // if (out_zone_offset_present)
 // *out_zone_offset_present = true;
 } else if (s[length - 5] == '+' || s[length - 5] == '-') {
 // [+-]HHMM
 length -= 5;
 if (ARROW_PREDICT_FALSE(!arrow::internal::detail::ParseHHMM(
 s + length + 1, &zone_offset))) {
 return false;
 }
 if (s[length] == '+')
 zone_offset *= -1;
 // if (out_zone_offset_present)
 // *out_zone_offset_present = true;
 } else if ((s[length - 6] == '+' || s[length - 6] == '-') &&
 (s[length - 3] == ':')) {
 // [+-]HH:MM
 length -= 6;
 if (ARROW_PREDICT_FALSE(!arrow::internal::detail::ParseHH_MM(
 s + length + 1, &zone_offset))) {
 return false;
 }
 if (s[length] == '+')
 zone_offset *= -1;
 // if (out_zone_offset_present)
 // *out_zone_offset_present = true;
 }
 
 seconds_type seconds_since_midnight;
 switch (length) {
 case 13: // YYYY-MM-DD[ T]hh
 if (ARROW_PREDICT_FALSE(!arrow::internal::detail::ParseHH(
 s + 11, &seconds_since_midnight))) {
 return false;
 }
 break;
 case 16: // YYYY-MM-DD[ T]hh:mm
 if (ARROW_PREDICT_FALSE(!arrow::internal::detail::ParseHH_MM(
 s + 11, &seconds_since_midnight))) {
 return false;
 }
 break;
 case 19: // YYYY-MM-DD[ T]hh:mm:ss
 case 21: // YYYY-MM-DD[ T]hh:mm:ss.s
 case 22: // YYYY-MM-DD[ T]hh:mm:ss.ss
 case 23: // YYYY-MM-DD[ T]hh:mm:ss.sss
 case 24: // YYYY-MM-DD[ T]hh:mm:ss.ssss
 case 25: // YYYY-MM-DD[ T]hh:mm:ss.sssss
 case 26: // YYYY-MM-DD[ T]hh:mm:ss.ssssss
 case 27: // YYYY-MM-DD[ T]hh:mm:ss.sssssss
 case 28: // YYYY-MM-DD[ T]hh:mm:ss.ssssssss
 case 29: // YYYY-MM-DD[ T]hh:mm:ss.sssssssss
 if (ARROW_PREDICT_FALSE(!arrow::internal::detail::ParseHH_MM_SS(
 s + 11, &seconds_since_midnight))) {
 return false;
 }
 break;
 default:
 LOG(ERROR) << "unsupported length: " << length;
 return false;
 }
 
 seconds_since_epoch += seconds_since_midnight;
 seconds_since_epoch += zone_offset;
 
 if (length <= 19) {
 *out =
 arrow::util::CastSecondsToUnit(out_unit, seconds_since_epoch.count());
 return true;
 }
 
 if (ARROW_PREDICT_FALSE(s[19] != '.')) {
 return false;
 }
 
 uint32_t subseconds = 0;
 if (ARROW_PREDICT_FALSE(!arrow::internal::detail::ParseSubSeconds(
 s + 20, length - 20, out_unit, &subseconds))) {
 return false;
 }
 
 *out =
 arrow::util::CastSecondsToUnit(out_unit, seconds_since_epoch.count()) +
 subseconds;
 return true;
 }
 
 const char* kind() const override { return "LDBC timestamp parser"; }
 
 const char* format() const override { return "EmptyFormat"; }
};
 
class LDBCLongDateParser : public arrow::TimestampParser {
 public:
 using seconds_type = std::chrono::duration<arrow::TimestampType::c_type>;
 LDBCLongDateParser() = default;
 
 ~LDBCLongDateParser() override {}
 
 bool operator()(const char* s, size_t length, arrow::TimeUnit::type out_unit,
 int64_t* out,
 bool* out_zone_offset_present = NULLPTR) const override {
 uint64_t seconds;
 // convert (s, s + length - 4) to seconds_since_epoch
 if (ARROW_PREDICT_FALSE(
 !arrow::internal::ParseUnsigned(s, length - 3, &seconds))) {
 return false;
 }
 
 uint32_t subseconds = 0;
 if (ARROW_PREDICT_FALSE(!arrow::internal::detail::ParseSubSeconds(
 s + length - 3, 3, out_unit, &subseconds))) {
 return false;
 }
 
 *out = arrow::util::CastSecondsToUnit(out_unit, seconds) + subseconds;
 return true;
 }
 
 const char* kind() const override { return "LDBC timestamp parser"; }
 
 const char* format() const override { return "LongDateFormat"; }
};
 
// convert c++ type to arrow type. support other types likes emptyType, Date
template <typename T>
struct TypeConverter;
 
template <>
struct TypeConverter<bool> {
 static PropertyType property_type() { return PropertyType::kBool; }
 using ArrowType = arrow::BooleanType;
 using ArrowArrayType = arrow::BooleanArray;
 static std::shared_ptr<arrow::DataType> ArrowTypeValue() {
 return arrow::boolean();
 }
};
 
template <>
struct TypeConverter<int32_t> {
 static PropertyType property_type() { return PropertyType::kInt32; }
 using ArrowType = arrow::Int32Type;
 using ArrowArrayType = arrow::Int32Array;
 static std::shared_ptr<arrow::DataType> ArrowTypeValue() {
 return arrow::int32();
 }
};
 
template <>
struct TypeConverter<uint32_t> {
 static PropertyType property_type() { return PropertyType::kUInt32; }
 using ArrowType = arrow::UInt32Type;
 using ArrowArrayType = arrow::UInt32Array;
 static std::shared_ptr<arrow::DataType> ArrowTypeValue() {
 return arrow::uint32();
 }
};
 
template <>
struct TypeConverter<int64_t> {
 static PropertyType property_type() { return PropertyType::kInt64; }
 using ArrowType = arrow::Int64Type;
 using ArrowArrayType = arrow::Int64Array;
 static std::shared_ptr<arrow::DataType> ArrowTypeValue() {
 return arrow::int64();
 }
};
 
template <>
struct TypeConverter<uint64_t> {
 static PropertyType property_type() { return PropertyType::kUInt64; }
 using ArrowType = arrow::UInt64Type;
 using ArrowArrayType = arrow::UInt64Array;
 static std::shared_ptr<arrow::DataType> ArrowTypeValue() {
 return arrow::uint64();
 }
};
 
template <>
struct TypeConverter<double> {
 static PropertyType property_type() { return PropertyType::kDouble; }
 using ArrowType = arrow::DoubleType;
 using ArrowArrayType = arrow::DoubleArray;
 static std::shared_ptr<arrow::DataType> ArrowTypeValue() {
 return arrow::float64();
 }
};
 
template <>
struct TypeConverter<float> {
 static PropertyType property_type() { return PropertyType::kFloat; }
 using ArrowType = arrow::FloatType;
 using ArrowArrayType = arrow::FloatArray;
 static std::shared_ptr<arrow::DataType> ArrowTypeValue() {
 return arrow::float32();
 }
};
template <>
struct TypeConverter<std::string> {
 static PropertyType property_type() { return PropertyType::kStringView; }
 using ArrowType = arrow::LargeStringType;
 using ArrowArrayType = arrow::LargeStringArray;
 static std::shared_ptr<arrow::DataType> ArrowTypeValue() {
 return arrow::large_utf8();
 }
};
 
template <>
struct TypeConverter<std::string_view> {
 static PropertyType property_type() { return PropertyType::kStringView; }
 using ArrowType = arrow::LargeStringType;
 using ArrowArrayType = arrow::LargeStringArray;
 static std::shared_ptr<arrow::DataType> ArrowTypeValue() {
 return arrow::large_utf8();
 }
};
 
template <>
struct TypeConverter<Date> {
 static PropertyType property_type() { return PropertyType::kDate; }
 using ArrowType = arrow::TimestampType;
 using ArrowArrayType = arrow::TimestampArray;
 static std::shared_ptr<arrow::DataType> ArrowTypeValue() {
 return arrow::timestamp(arrow::TimeUnit::MILLI);
 }
};
 
template <>
struct TypeConverter<Day> {
 static PropertyType property_type() { return PropertyType::kDay; }
 using ArrowType = arrow::TimestampType;
 using ArrowArrayType = arrow::TimestampArray;
 static std::shared_ptr<arrow::DataType> ArrowTypeValue() {
 return arrow::timestamp(arrow::TimeUnit::MILLI);
 }
};
 
std::shared_ptr<arrow::DataType> PropertyTypeToArrowType(PropertyType type);
} // namespace gs
 
#endif // UTILS_ARROW_UTILS_H_