Can read date32 as PostgreSQL date from arrow (#1046)

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