arrow_impl.h 1.9 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657
  1. #pragma once
  2. #include <arrow/array.h>
  3. #include <arrow/array/builder_binary.h>
  4. #include <yql/essentials/parser/pg_wrapper/interface/arrow.h>
  5. extern "C" {
  6. #include "utils/numeric.h"
  7. }
  8. namespace NYql {
  9. Numeric Uint64ToPgNumeric(ui64 value);
  10. Numeric DecimalToPgNumeric(const NUdf::TUnboxedValuePod& value, ui8 precision, ui8 scale);
  11. Numeric DyNumberToPgNumeric(const NUdf::TUnboxedValuePod& value);
  12. Numeric PgFloatToNumeric(double item, ui64 scale, int digits);
  13. Numeric PgDecimal128ToNumeric(arrow::Decimal128 val, int32_t precision, int32_t scale, Numeric high_bits_mul);
  14. TColumnConverter BuildPgColumnConverter(const std::shared_ptr<arrow::DataType>& originalType, NKikimr::NMiniKQL::TPgType* targetType);
  15. template<typename T>
  16. std::shared_ptr<arrow::Array> PgConvertNumeric(const std::shared_ptr<arrow::Array>& value) {
  17. TArenaMemoryContext arena;
  18. const auto& data = value->data();
  19. size_t length = data->length;
  20. arrow::BinaryBuilder builder;
  21. auto input = data->GetValues<T>(1);
  22. for (size_t i = 0; i < length; ++i) {
  23. if (value->IsNull(i)) {
  24. ARROW_OK(builder.AppendNull());
  25. continue;
  26. }
  27. T item = input[i];
  28. Numeric v;
  29. if constexpr(std::is_same_v<T, double>) {
  30. v = PgFloatToNumeric(item, 1000000000000LL, 12);
  31. } else if constexpr(std::is_same_v<T, float>) {
  32. v = PgFloatToNumeric(item, 1000000LL, 6);
  33. } else {
  34. v = int64_to_numeric(item);
  35. }
  36. auto datum = NumericGetDatum(v);
  37. auto ptr = (char*)datum;
  38. auto len = GetFullVarSize((const text*)datum);
  39. NUdf::ZeroMemoryContext(ptr);
  40. ARROW_OK(builder.Append(ptr - sizeof(void*), len + sizeof(void*)));
  41. }
  42. std::shared_ptr<arrow::BinaryArray> ret;
  43. ARROW_OK(builder.Finish(&ret));
  44. return ret;
  45. }
  46. std::shared_ptr<arrow::Array> PgDecimal128ConvertNumeric(const std::shared_ptr<arrow::Array>& value, int32_t precision, int32_t scale);
  47. }