diff options
author | Alexey Ozeritskiy <aozeritsky@ydb.tech> | 2024-01-16 17:34:24 +0100 |
---|---|---|
committer | GitHub <noreply@github.com> | 2024-01-16 17:34:24 +0100 |
commit | 60c669702c975cea03576d9f6ee7f371c9b3a24e (patch) | |
tree | 8e6d0e20f0ce049b0a547ce385a3a8e932207ea6 | |
parent | c01d0002fc70f9946d70e5280d050fb1e9121f48 (diff) | |
download | ydb-60c669702c975cea03576d9f6ee7f371c9b3a24e.tar.gz |
Can read date32 as PostgreSQL date from arrow (#1046)
-rw-r--r-- | ydb/library/yql/parser/pg_wrapper/arrow.cpp | 4 | ||||
-rw-r--r-- | ydb/library/yql/parser/pg_wrapper/arrow_impl.h | 2 | ||||
-rw-r--r-- | ydb/library/yql/parser/pg_wrapper/ut/arrow_ut.cpp | 42 |
3 files changed, 48 insertions, 0 deletions
diff --git a/ydb/library/yql/parser/pg_wrapper/arrow.cpp b/ydb/library/yql/parser/pg_wrapper/arrow.cpp index 5a8559ff8f..0bb21f5706 100644 --- a/ydb/library/yql/parser/pg_wrapper/arrow.cpp +++ b/ydb/library/yql/parser/pg_wrapper/arrow.cpp @@ -276,6 +276,10 @@ TColumnConverter BuildPgColumnConverter(const std::shared_ptr<arrow::DataType>& return [](const std::shared_ptr<arrow::Array>& value) { return PgConvertFixed<ui16>(value, [](auto value){ return MakePgDateFromUint16(value); }); }; + } else if (originalType->Equals(arrow::date32())) { + return [](const std::shared_ptr<arrow::Array>& value) { + return PgConvertFixed<i32>(value, [](auto value){ return MakePgDateFromUint16(value); }); + }; } else { return {}; } diff --git a/ydb/library/yql/parser/pg_wrapper/arrow_impl.h b/ydb/library/yql/parser/pg_wrapper/arrow_impl.h index 7574aada42..62daeea279 100644 --- a/ydb/library/yql/parser/pg_wrapper/arrow_impl.h +++ b/ydb/library/yql/parser/pg_wrapper/arrow_impl.h @@ -2,6 +2,7 @@ #include <arrow/array.h> #include <arrow/array/builder_binary.h> +#include <ydb/library/yql/parser/pg_wrapper/interface/arrow.h> extern "C" { #include "utils/numeric.h" @@ -10,6 +11,7 @@ extern "C" { namespace NYql { Numeric PgFloatToNumeric(double item, ui64 scale, int digits); +TColumnConverter BuildPgColumnConverter(const std::shared_ptr<arrow::DataType>& originalType, NKikimr::NMiniKQL::TPgType* targetType); template<typename T> std::shared_ptr<arrow::Array> PgConvertNumeric(const std::shared_ptr<arrow::Array>& value) { diff --git a/ydb/library/yql/parser/pg_wrapper/ut/arrow_ut.cpp b/ydb/library/yql/parser/pg_wrapper/ut/arrow_ut.cpp index c09359f405..28caad945d 100644 --- a/ydb/library/yql/parser/pg_wrapper/ut/arrow_ut.cpp +++ b/ydb/library/yql/parser/pg_wrapper/ut/arrow_ut.cpp @@ -106,6 +106,48 @@ Y_UNIT_TEST(PgConvertNumericInt) { } } +Y_UNIT_TEST(PgConvertDate32Date) { + TArenaMemoryContext arena; + + arrow::Date32Builder builder; + builder.Append(10227); + builder.AppendNull(); + builder.Append(11323); + builder.Append(10227); + builder.Append(10958); + builder.Append(11688); + + std::shared_ptr<arrow::Array> array; + builder.Finish(&array); + + NKikimr::NMiniKQL::TScopedAlloc alloc(__LOCATION__); + NKikimr::NMiniKQL::TTypeEnvironment typeEnv(alloc); + auto* targetType = NKikimr::NMiniKQL::TPgType::Create(DATEOID, typeEnv); + + auto converter = BuildPgColumnConverter(std::shared_ptr<arrow::DataType>(new arrow::Date32Type), targetType); + auto result = converter(array); + const auto& data = result->data(); + UNIT_ASSERT_VALUES_EQUAL(result->length(), 6); + + const char* expected[] = { + "1998-01-01", nullptr, "2001-01-01", "1998-01-01", "2000-01-02", "2002-01-01" + }; + + NUdf::TFixedSizeBlockReader<ui64, true> reader; + for (int i = 0; i < 6; i++) { + if (result->IsNull(i)) { + UNIT_ASSERT(expected[i] == nullptr); + } else { + auto item = reader.GetItem(*data, i).As<Datum>(); + UNIT_ASSERT(expected[i] != nullptr); + UNIT_ASSERT_VALUES_EQUAL( + TString(DatumGetCString(DirectFunctionCall1(date_out, item))), + expected[i] + ); + } + } +} + } // Y_UNIT_TEST_SUITE(TArrowUtilsTests) } // namespace NYql |