aboutsummaryrefslogtreecommitdiffstats
diff options
context:
space:
mode:
authorAlexey Ozeritskiy <aozeritsky@ydb.tech>2024-01-16 17:34:24 +0100
committerGitHub <noreply@github.com>2024-01-16 17:34:24 +0100
commit60c669702c975cea03576d9f6ee7f371c9b3a24e (patch)
tree8e6d0e20f0ce049b0a547ce385a3a8e932207ea6
parentc01d0002fc70f9946d70e5280d050fb1e9121f48 (diff)
downloadydb-60c669702c975cea03576d9f6ee7f371c9b3a24e.tar.gz
Can read date32 as PostgreSQL date from arrow (#1046)
-rw-r--r--ydb/library/yql/parser/pg_wrapper/arrow.cpp4
-rw-r--r--ydb/library/yql/parser/pg_wrapper/arrow_impl.h2
-rw-r--r--ydb/library/yql/parser/pg_wrapper/ut/arrow_ut.cpp42
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