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 2ff79b30e739..a726fe16f9fa 100644 --- a/ydb/library/yql/parser/pg_wrapper/ut/arrow_ut.cpp +++ b/ydb/library/yql/parser/pg_wrapper/ut/arrow_ut.cpp @@ -12,18 +12,24 @@ extern "C" { namespace { -void checkResult(const char ** expected, std::shared_ptr data) { +template +void checkResult(const char ** expected, auto result, NYql::NUdf::IBlockReader* reader, auto out_fun) { + const auto& data = result->data(); - NYql::NUdf::TStringBlockReader reader; for (int i = 0; i < data->length; i++) { - auto item = reader.GetItem(*data, i); - if (!item) { + if (result->IsNull(i)) { UNIT_ASSERT(expected[i] == nullptr); } else { - const char* addr = item.AsStringRef().Data() + sizeof(void*); UNIT_ASSERT(expected[i] != nullptr); + + Datum item; + if constexpr (IsFixedSizeReader) { + item = reader->GetItem(*data, i).template As(); + } else { + item = Datum(reader->GetItem(*data, i).AsStringRef().Data() + sizeof(void*)); + } UNIT_ASSERT_VALUES_EQUAL( - TString(DatumGetCString(DirectFunctionCall1(numeric_out, (Datum)addr))), + TString(DatumGetCString(DirectFunctionCall1(out_fun, item))), expected[i] ); } @@ -70,15 +76,13 @@ Y_UNIT_TEST(PgConvertNumericDouble) { builder.Finish(&array); auto result = PgConvertNumeric(array); - const auto& data = result->data(); - - const char* expected[] = { "1.1", "31.37", nullptr, "-1.337", "0" }; - checkResult(expected, data); + NYql::NUdf::TStringBlockReader reader; + checkResult(expected, result, &reader, numeric_out); } Y_UNIT_TEST(PgConvertNumericInt) { @@ -101,7 +105,8 @@ Y_UNIT_TEST(PgConvertNumericInt) { "11", "3137", nullptr, "-1337", "0" }; - checkResult(expected, data); + NYql::NUdf::TStringBlockReader reader; + checkResult(expected, result, &reader, numeric_out); } Y_UNIT_TEST(PgConvertDate32Date) { @@ -124,7 +129,6 @@ Y_UNIT_TEST(PgConvertDate32Date) { auto converter = BuildPgColumnConverter(std::shared_ptr(new arrow::Date32Type), targetType); auto result = converter(array); - const auto& data = result->data(); UNIT_ASSERT_VALUES_EQUAL(result->length(), 6); const char* expected[] = { @@ -132,18 +136,7 @@ Y_UNIT_TEST(PgConvertDate32Date) { }; NUdf::TFixedSizeBlockReader 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(); - UNIT_ASSERT(expected[i] != nullptr); - UNIT_ASSERT_VALUES_EQUAL( - TString(DatumGetCString(DirectFunctionCall1(date_out, item))), - expected[i] - ); - } - } + checkResult(expected, result, &reader, date_out); } } // Y_UNIT_TEST_SUITE(TArrowUtilsTests)