From e1cd64df2f5bf9ed2148cb2cc81bb9686751f102 Mon Sep 17 00:00:00 2001 From: mrlolthe1st Date: Wed, 25 Dec 2024 16:36:31 +0300 Subject: YQL-18548: BlockReader PG types + refactor fix fix commit_hash:09713e813e389b07077111daed2ec31b59047bca --- yql/essentials/parser/pg_wrapper/arrow.cpp | 553 ++++++++++++++++++++- yql/essentials/parser/pg_wrapper/interface/arrow.h | 10 +- yql/essentials/parser/pg_wrapper/interface/ya.make | 1 + yql/essentials/providers/common/codec/ya.make | 1 + .../codec/yt_arrow_converter_interface/ya.make | 10 + .../yt_arrow_converter.h | 13 + .../yt_arrow_converter_details.h | 100 ++++ yql/essentials/sql/pg_dummy/pg_sql_dummy.cpp | 10 + yt/yql/providers/yt/codec/README.md | 4 + yt/yql/providers/yt/codec/ya.make | 2 + yt/yql/providers/yt/codec/yt_arrow_converter.cpp | 247 ++++++--- yt/yql/providers/yt/codec/yt_arrow_converter.h | 21 +- .../yt/codec/yt_arrow_converter_details.h | 129 ----- yt/yql/providers/yt/codec/yt_codec_io.cpp | 2 + .../yt/comp_nodes/dq/dq_yt_block_reader.cpp | 3 + 15 files changed, 872 insertions(+), 234 deletions(-) create mode 100644 yql/essentials/providers/common/codec/yt_arrow_converter_interface/ya.make create mode 100644 yql/essentials/providers/common/codec/yt_arrow_converter_interface/yt_arrow_converter.h create mode 100644 yql/essentials/providers/common/codec/yt_arrow_converter_interface/yt_arrow_converter_details.h create mode 100644 yt/yql/providers/yt/codec/README.md delete mode 100644 yt/yql/providers/yt/codec/yt_arrow_converter_details.h diff --git a/yql/essentials/parser/pg_wrapper/arrow.cpp b/yql/essentials/parser/pg_wrapper/arrow.cpp index 9bfe088b014..8c966e4f452 100644 --- a/yql/essentials/parser/pg_wrapper/arrow.cpp +++ b/yql/essentials/parser/pg_wrapper/arrow.cpp @@ -8,6 +8,7 @@ #include #include #include + #include #include @@ -83,7 +84,7 @@ std::shared_ptr PgConvertBool(const std::shared_ptr& size_t length = data->length; NUdf::TFixedSizeArrayBuilder builder(NKikimr::NMiniKQL::TTypeInfoHelper(), arrow::uint64(), *arrow::default_memory_pool(), length); auto input = data->GetValues(1, 0); - builder.UnsafeReserve(length); + builder.UnsafeReserve(length); auto output = builder.MutableData(); for (size_t i = 0; i < length; ++i) { auto fullIndex = i + data->offset; @@ -179,7 +180,7 @@ Numeric Uint64ToPgNumeric(ui64 value) { Numeric DecimalToPgNumeric(const NUdf::TUnboxedValuePod& value, ui8 precision, ui8 scale) { const auto str = NYql::NDecimal::ToString(value.GetInt128(), precision, scale); Y_ENSURE(str); - return (Numeric)DirectFunctionCall3Coll(numeric_in, DEFAULT_COLLATION_OID, + return (Numeric)DirectFunctionCall3Coll(numeric_in, DEFAULT_COLLATION_OID, PointerGetDatum(str), Int32GetDatum(0), Int32GetDatum(-1)); } @@ -229,12 +230,12 @@ std::shared_ptr PgDecimal128ConvertNumeric(const std::shared_ptr ret; @@ -254,7 +255,7 @@ Numeric PgDecimal128ToNumeric(arrow::Decimal128 value, int32_t precision, int32_ Numeric low_bits_res = int64_div_fast_to_numeric(low_bits, scale); Numeric high_bits_res = numeric_mul_opt_error(int64_div_fast_to_numeric(high_bits, scale), high_bits_mul, &error); MKQL_ENSURE(error == false, "Bad numeric multiplication."); - + Numeric res = numeric_add_opt_error(high_bits_res, low_bits_res, &error); MKQL_ENSURE(error == false, "Bad numeric addition."); @@ -335,7 +336,7 @@ TColumnConverter BuildPgColumnConverter(const std::shared_ptr& } case INT4OID: { return BuildPgFixedColumnConverter(originalType, [](auto value){ return Int32GetDatum(value); }); - } + } case INT8OID: { return BuildPgFixedColumnConverter(originalType, [](auto value){ return Int64GetDatum(value); }); } @@ -387,9 +388,547 @@ TColumnConverter BuildPgColumnConverter(const std::shared_ptr& } else { return {}; } - } + } } return {}; } +class IYsonBlockReaderForPg : public IYsonComplexTypeReader { +public: + virtual NUdf::TBlockItem GetNotNull(TYsonBuffer&) = 0; + NUdf::TBlockItem GetNullableItem(TYsonBuffer& buf) { + char prev = buf.Current(); + if (prev == NYson::NDetail::EntitySymbol) { + buf.Next(); + return NUdf::TBlockItem(); + } + if (prev == NYson::NDetail::BeginListSymbol) { + buf.Next(); + YQL_ENSURE(buf.Current() == NYson::NDetail::EndListSymbol); + buf.Next(); + return NUdf::TBlockItem(); + } + return GetNotNull(buf); + } +}; + + +NUdf::TBlockItem BlockItemFromDatum(Datum datum, const NPg::TTypeDesc& desc, std::vector& tmp) { + if (desc.PassByValue) { + return NUdf::TBlockItem((ui64)datum); + } + auto typeLen = desc.TypeLen; + ui32 len; + if (typeLen == -1) { + len = GetFullVarSize((const text*)datum); + } else if (typeLen == -2) { + len = 1 + strlen((const char*)datum); + } else { + len = typeLen; + } + auto objlen = len; + len += sizeof(void*); + len = AlignUp(len, 8); + tmp.resize(len); + *(uint64_t*)tmp.data() = 0; + memcpy(tmp.data() + sizeof(void*), (const char*) datum, objlen); + return NUdf::TBlockItem(std::string_view(tmp.data(), len)); +} + +NUdf::TBlockItem PgBlockItemFromNativeBinary(const TStringBuf binary, ui32 pgTypeId, std::vector& tmp) { + NKikimr::NMiniKQL::TPAllocScope call; + StringInfoData stringInfo; + stringInfo.data = (char*)binary.Data(); + stringInfo.len = binary.Size(); + stringInfo.maxlen = binary.Size(); + stringInfo.cursor = 0; + + const auto& typeInfo = NPg::LookupType(pgTypeId); + auto typeIOParam = MakeTypeIOParam(typeInfo); + auto receiveFuncId = typeInfo.ReceiveFuncId; + if (typeInfo.TypeId == typeInfo.ArrayTypeId) { + receiveFuncId = NPg::LookupProc("array_recv", { 0,0,0 }).ProcId; + } + + { + FmgrInfo finfo; + Zero(finfo); + Y_ENSURE(receiveFuncId); + fmgr_info(receiveFuncId, &finfo); + Y_ENSURE(!finfo.fn_retset); + Y_ENSURE(finfo.fn_addr); + Y_ENSURE(finfo.fn_nargs >= 1 && finfo.fn_nargs <= 3); + LOCAL_FCINFO(callInfo, 3); + Zero(*callInfo); + callInfo->flinfo = &finfo; + callInfo->nargs = 3; + callInfo->fncollation = DEFAULT_COLLATION_OID; + callInfo->isnull = false; + callInfo->args[0] = { (Datum)&stringInfo, false }; + callInfo->args[1] = { ObjectIdGetDatum(typeIOParam), false }; + callInfo->args[2] = { Int32GetDatum(-1), false }; + + auto x = finfo.fn_addr(callInfo); + Y_ENSURE(!callInfo->isnull); + if (stringInfo.cursor != stringInfo.len) { + TStringBuilder errMsg; + errMsg << "Not all data has been consumed by 'recv' function: " << NPg::LookupProc(receiveFuncId).Name << ", data size: " << stringInfo.len << ", consumed size: " << stringInfo.cursor; + UdfTerminate(errMsg.c_str()); + } + return BlockItemFromDatum(x, typeInfo, tmp); + } +} + +template +constexpr Datum FixedToDatum(T v) { + if constexpr (std::is_same_v) { + return BoolGetDatum(v); + } else if constexpr (std::is_same_v) { + return Int16GetDatum(v); + } else if constexpr (std::is_same_v) { + return Int32GetDatum(v); + } else if constexpr (std::is_same_v) { + return Int64GetDatum(v); + } else if constexpr (std::is_same_v) { + return Float4GetDatum(v); + } else if constexpr (std::is_same_v) { + return Float8GetDatum(v); + } +} + +template +class TPgYsonFixedConverter final : public IYsonBlockReaderForPg { +public: + NUdf::TBlockItem GetItem(TYsonBuffer& buf) override final { + return this->GetNullableItem(buf); + } + + NUdf::TBlockItem GetNotNull(TYsonBuffer& buf) override final { + Datum val; + if constexpr (std::is_same_v) { + Y_ENSURE(buf.Current() == NYson::NDetail::FalseMarker || buf.Current() == NYson::NDetail::TrueMarker); + val = FixedToDatum(buf.Current() == NYson::NDetail::TrueMarker); + buf.Next(); + } else if constexpr (std::is_integral_v) { + if constexpr (std::is_signed_v) { + Y_ENSURE(buf.Current() == NYson::NDetail::Int64Marker); + buf.Next(); + val = FixedToDatum(buf.ReadVarI64()); + } else { + Y_ENSURE(buf.Current() == NYson::NDetail::Uint64Marker); + buf.Next(); + val = FixedToDatum(buf.ReadVarUI64()); + } + } else { + Y_ENSURE(buf.Current() == NYson::NDetail::DoubleMarker); + buf.Next(); + val = FixedToDatum(buf.NextDouble()); + } + return NUdf::TBlockItem(val); + } +}; + +template +class TPgYsonStringConverter final : public IYsonBlockReaderForPg { +public: + + TPgYsonStringConverter(i32 typeLen) : TypeLen_(typeLen) { + if (typeLen == -2) { + YQL_ENSURE(IsCString && !FixedLength); + } else if (typeLen == -1) { + YQL_ENSURE(!IsCString && !FixedLength); + } else { + YQL_ENSURE(typeLen >= 0 && FixedLength); + } + } + + NUdf::TBlockItem GetItem(TYsonBuffer& buf) override final { + return this->GetNullableItem(buf); + } + + NUdf::TBlockItem GetNotNull(TYsonBuffer& buf) override final { + Y_ENSURE(buf.Current() == NYson::NDetail::StringMarker); + buf.Next(); + const i32 originalLen = buf.ReadVarI32(); + auto res = buf.Data(); + buf.Skip(originalLen); + + ui32 len; + if constexpr (IsCString) { + len = 1 + originalLen + sizeof(void*); + } else if constexpr (FixedLength) { + len = TypeLen_ + sizeof(void*); + } else { + len = VARHDRSZ + originalLen + sizeof(void*); + } + + if (Tmp_.capacity() < len) { + Tmp_.reserve(Max(len, Tmp_.capacity() * 2)); + } + len = AlignUp(len, 8); + Tmp_.resize(len); + if constexpr (IsCString) { + memcpy(Tmp_.data() + sizeof(void*), res, originalLen); + } else if constexpr (FixedLength) { + memcpy(Tmp_.data() + sizeof(void*), res, originalLen); + } else { + memcpy(Tmp_.data() + VARHDRSZ + sizeof(void*), res, originalLen); + UpdateCleanVarSize((text*)(Tmp_.data() + sizeof(void*)), originalLen); + } + return NUdf::TBlockItem(NUdf::TStringRef(Tmp_.data(), len)); + } +private: + std::vector Tmp_; + i32 TypeLen_; +}; + + +class TPgYsonOtherConverter : public IYsonBlockReaderForPg { +public: + TPgYsonOtherConverter(Oid typeId) : TypeId_(typeId) {} + NUdf::TBlockItem GetItem(TYsonBuffer& buf) override final { + return this->GetNullableItem(buf); + } + + NUdf::TBlockItem GetNotNull(TYsonBuffer& buf) override final { + if (buf.Current() != NYson::NDetail::StringMarker) { + Y_ENSURE(buf.Current() == NYson::NDetail::StringMarker); + } + buf.Next(); + const i32 len = buf.ReadVarI32(); + auto ptr = buf.Data(); + buf.Skip(len); + return PgBlockItemFromNativeBinary(TStringBuf(ptr, len), TypeId_, Tmp_); + } +private: + Oid TypeId_; + std::vector Tmp_; +}; + + +template +class TPgTopLevelFixedConverter : public IYtColumnConverter { +public: + using Fn = Datum(*)(const T&); + TPgTopLevelFixedConverter(std::unique_ptr&& builder) : Builder_(std::move(builder)) {} + + arrow::Datum Convert(std::shared_ptr data) override final { + if (arrow::Type::DICTIONARY == data->type->id()) { + auto valType = static_cast(*data->type).value_type(); + Y_ENSURE(Expected == valType->id()); + return ConvertDict(data); + } else { + Y_ENSURE(Expected == data->type->id()); + return ConvertNonDict(data); + } + } + + arrow::Datum ConvertNonDict(std::shared_ptr data) { + ArrType arr(data); + if (arr.null_count()) { + for (i64 i = 0; i < data->length; ++i) { + if (arr.IsNull(i)) { + Builder_->Add(NUdf::TBlockItem{}); + } else { + Builder_->Add(NUdf::TBlockItem(FixedToDatum(arr.Value(i)))); + } + } + } else { + for (i64 i = 0; i < data->length; ++i) { + Builder_->Add(NUdf::TBlockItem(FixedToDatum(arr.Value(i)))); + } + } + return Builder_->Build(false); + } + + arrow::Datum ConvertDict(std::shared_ptr data) { + arrow::DictionaryArray dict(data); + auto values = data->dictionary->GetValues(0); + auto indices = dict.indices()->data()->GetValues(1); + if (dict.null_count()) { + for (i64 i = 0; i < data->length; ++i) { + if (dict.IsNull(i)) { + Builder_->Add(NUdf::TBlockItem{}); + } else { + Builder_->Add(NUdf::TBlockItem(FixedToDatum(values[indices[i]]))); + } + } + } else { + for (i64 i = 0; i < data->length; ++i) { + Builder_->Add(NUdf::TBlockItem(FixedToDatum(values[indices[i]]))); + } + } + return Builder_->Build(false); + } +private: + std::unique_ptr Builder_; +}; + +template +class TPgTopLevelStringConverter : public IYtColumnConverter { +public: + TPgTopLevelStringConverter(std::unique_ptr&& builder, i32 typeLen) : Builder_(std::move(builder)), TypeLen_(typeLen) { + if (typeLen == -2) { + YQL_ENSURE(IsCString && !FixedLength); + } else if (typeLen == -1) { + YQL_ENSURE(!IsCString && !FixedLength); + } else { + YQL_ENSURE(typeLen >= 0 && FixedLength); + } + } + + constexpr NUdf::TBlockItem ConvertOnce(const uint8_t* res, size_t originalLen) { + ui32 len; + if constexpr (IsCString) { + len = 1 + originalLen + sizeof(void*); + } else if constexpr (FixedLength) { + len = TypeLen_ + sizeof(void*); + } else { + len = VARHDRSZ + originalLen + sizeof(void*); + } + + if (Tmp_.capacity() < len) { + Tmp_.reserve(Max(len, Tmp_.capacity() * 2)); + } + len = AlignUp(len, 8); + Tmp_.resize(len); + if constexpr (IsCString) { + memcpy(Tmp_.data() + sizeof(void*), res, originalLen); + } else if constexpr (FixedLength) { + memcpy(Tmp_.data() + sizeof(void*), res, originalLen); + } else { + memcpy(Tmp_.data() + VARHDRSZ + sizeof(void*), res, originalLen); + UpdateCleanVarSize((text*)(Tmp_.data() + sizeof(void*)), originalLen); + } + return NUdf::TBlockItem(NUdf::TStringRef(Tmp_.data(), len)); + } + + arrow::Datum Convert(std::shared_ptr data) override final { + if (arrow::Type::DICTIONARY == data->type->id()) { + auto valType = static_cast(*data->type).value_type(); + Y_ENSURE(arrow::Type::BINARY == valType->id() || arrow::Type::STRING == valType->id()); + return ConvertDict(data); + } else { + if (arrow::Type::STRING == data->type->id()) { + auto res = arrow::compute::Cast(data, std::make_shared()); + Y_ENSURE(res.ok()); + data = res->array(); + } + Y_ENSURE(arrow::Type::BINARY == data->type->id()); + return ConvertNonDict(data); + } + } + + arrow::Datum ConvertNonDict(std::shared_ptr data) { + arrow::BinaryArray arr(data); + if (arr.null_count()) { + for (i64 i = 0; i < data->length; ++i) { + if (arr.IsNull(i)) { + Builder_->Add(NUdf::TBlockItem{}); + } else { + i32 len; + auto res = arr.GetValue(i, &len); + Builder_->Add(ConvertOnce(res, len)); + } + } + } else { + for (i64 i = 0; i < data->length; ++i) { + i32 len; + auto res = arr.GetValue(i, &len); + Builder_->Add(ConvertOnce(res, len)); + } + } + return Builder_->Build(false); + } + + arrow::Datum ConvertDict(std::shared_ptr data) { + arrow::DictionaryArray dict(data); + if (arrow::Type::STRING == data->dictionary->type->id()) { + auto res = arrow::compute::Cast(data->dictionary, std::make_shared()); + Y_ENSURE(res.ok()); + data->dictionary = res->array(); + } + arrow::BinaryArray arr(data->dictionary); + auto indices = dict.indices()->data()->GetValues(1); + if (dict.null_count()) { + for (i64 i = 0; i < data->length; ++i) { + if (dict.IsNull(i)) { + Builder_->Add(NUdf::TBlockItem{}); + } else { + i32 len; + auto res = arr.GetValue(indices[i], &len); + Builder_->Add(NUdf::TBlockItem(ConvertOnce(res, len))); + } + } + } else { + for (i64 i = 0; i < data->length; ++i) { + i32 len; + auto res = arr.GetValue(indices[i], &len); + Builder_->Add(NUdf::TBlockItem(ConvertOnce(res, len))); + } + } + return Builder_->Build(false); + } +private: + std::unique_ptr Builder_; + std::vector Tmp_; + i32 TypeLen_; +}; + +class TPgTopLevelOtherConverter : public IYtColumnConverter { +public: + TPgTopLevelOtherConverter(std::unique_ptr&& builder, Oid typeId) : Builder_(std::move(builder)), TypeId_(typeId) {} + + inline NUdf::TBlockItem ConvertOnce(const uint8_t* res, size_t len) { + return PgBlockItemFromNativeBinary(TStringBuf(reinterpret_cast(res), len), TypeId_, Tmp_); + } + + arrow::Datum Convert(std::shared_ptr data) override final { + if (arrow::Type::DICTIONARY == data->type->id()) { + auto valType = static_cast(*data->type).value_type(); + Y_ENSURE(arrow::Type::BINARY == valType->id() || arrow::Type::STRING == valType->id()); + return ConvertDict(data); + } else { + Y_ENSURE(arrow::Type::BINARY == data->type->id() || arrow::Type::STRING == data->type->id()); + return ConvertNonDict(data); + } + } + + arrow::Datum ConvertNonDict(std::shared_ptr data) { + arrow::BinaryArray arr(data); + if (arr.null_count()) { + for (i64 i = 0; i < data->length; ++i) { + if (arr.IsNull(i)) { + Builder_->Add(NUdf::TBlockItem{}); + } else { + i32 len; + auto res = arr.GetValue(i, &len); + Builder_->Add(NUdf::TBlockItem(ConvertOnce(res, len))); + } + } + } else { + for (i64 i = 0; i < data->length; ++i) { + i32 len; + auto res = arr.GetValue(i, &len); + Builder_->Add(NUdf::TBlockItem(ConvertOnce(res, len))); + } + } + return Builder_->Build(false); + } + + arrow::Datum ConvertDict(std::shared_ptr data) { + arrow::DictionaryArray dict(data); + if (arrow::Type::STRING == data->dictionary->type->id()) { + auto res = arrow::compute::Cast(data->dictionary, std::make_shared()); + Y_ENSURE(res.ok()); + data->dictionary = res->array(); + } + arrow::BinaryArray arr(data->dictionary); + auto indices = dict.indices()->data()->GetValues(1); + if (dict.null_count()) { + for (i64 i = 0; i < data->length; ++i) { + if (dict.IsNull(i)) { + Builder_->Add(NUdf::TBlockItem{}); + } else { + i32 len; + auto res = arr.GetValue(indices[i], &len); + Builder_->Add(NUdf::TBlockItem(ConvertOnce(res, len))); + } + } + } else { + for (i64 i = 0; i < data->length; ++i) { + i32 len; + auto res = arr.GetValue(indices[i], &len); + Builder_->Add(NUdf::TBlockItem(ConvertOnce(res, len))); + } + } + return Builder_->Build(false); + } +private: + std::unique_ptr Builder_; + Oid TypeId_; + std::vector Tmp_; +}; + +std::unique_ptr BuildPgTopLevelColumnReader(std::unique_ptr&& builder, const NKikimr::NMiniKQL::TPgType* targetType) { + YQL_ENSURE(targetType); + + switch (targetType->GetTypeId()) { + case BOOLOID: { + return std::make_unique>(std::move(builder)); + } + case INT2OID: { + return std::make_unique>(std::move(builder)); + } + case INT4OID: { + return std::make_unique>(std::move(builder)); + } + case INT8OID: { + return std::make_unique>(std::move(builder)); + } + case FLOAT4OID: { + return std::make_unique>(std::move(builder)); + } + case FLOAT8OID: { + return std::make_unique>(std::move(builder)); + } + case BYTEAOID: + case VARCHAROID: + case TEXTOID: + case NAMEOID: + case CSTRINGOID: { + auto typeLen = NPg::LookupType(targetType->GetTypeId()).TypeLen; + if (typeLen == -2) { + return std::make_unique>(std::move(builder), typeLen); + } else if (typeLen == -1) { + return std::make_unique>(std::move(builder), typeLen); + } else { + return std::make_unique>(std::move(builder), typeLen); + } + } + default: + return std::make_unique(std::move(builder), targetType->GetTypeId()); + } +} + + +std::unique_ptr BuildPgYsonColumnReader(const NUdf::TPgTypeDescription& desc) { + switch (desc.TypeId) { + case BOOLOID: { + return std::make_unique>(); + } + case INT2OID: { + return std::make_unique>(); + } + case INT4OID: { + return std::make_unique>(); + } + case INT8OID: { + return std::make_unique>(); + } + case FLOAT4OID: { + return std::make_unique>(); + } + case FLOAT8OID: { + return std::make_unique>(); + } + case BYTEAOID: + case NAMEOID: + case VARCHAROID: + case TEXTOID: + case CSTRINGOID: { + auto typeLen = NPg::LookupType(desc.TypeId).TypeLen; + if (typeLen == -2) { + return std::make_unique>(typeLen); + } else if (typeLen == -1) { + return std::make_unique>(typeLen); + } else { + return std::make_unique>(typeLen); + } + } + default: + return std::make_unique(desc.TypeId); + } +} + } diff --git a/yql/essentials/parser/pg_wrapper/interface/arrow.h b/yql/essentials/parser/pg_wrapper/interface/arrow.h index f8e9c2a0622..f9aead6a553 100644 --- a/yql/essentials/parser/pg_wrapper/interface/arrow.h +++ b/yql/essentials/parser/pg_wrapper/interface/arrow.h @@ -1,8 +1,14 @@ #pragma once +#include +#include + #include #include +#include + #include + namespace NYql { arrow::Datum MakePgScalar(NKikimr::NMiniKQL::TPgType* type, const NKikimr::NUdf::TUnboxedValuePod& value, arrow::MemoryPool& pool); @@ -11,6 +17,8 @@ arrow::Datum MakePgScalar(NKikimr::NMiniKQL::TPgType* type, const NUdf::TBlockIt using TColumnConverter = std::function(const std::shared_ptr&)>; TColumnConverter BuildPgColumnConverter(const std::shared_ptr& originalType, NKikimr::NMiniKQL::TPgType* targetType); +std::unique_ptr BuildPgYsonColumnReader(const NUdf::TPgTypeDescription& desc); +std::unique_ptr BuildPgTopLevelColumnReader(std::unique_ptr&& builder, const NKikimr::NMiniKQL::TPgType* targetType); } // NYql namespace NKikimr { @@ -20,4 +28,4 @@ class IBlockAggregatorFactory; void RegisterPgBlockAggs(THashMap>& registry); } -} \ No newline at end of file +} diff --git a/yql/essentials/parser/pg_wrapper/interface/ya.make b/yql/essentials/parser/pg_wrapper/interface/ya.make index 92fe7ff1669..90f37bffcdf 100644 --- a/yql/essentials/parser/pg_wrapper/interface/ya.make +++ b/yql/essentials/parser/pg_wrapper/interface/ya.make @@ -20,6 +20,7 @@ PEERDIR( yql/essentials/public/udf/arrow yql/essentials/core/cbo library/cpp/disjoint_sets + yql/essentials/providers/common/codec/yt_arrow_converter_interface ) YQL_LAST_ABI_VERSION() diff --git a/yql/essentials/providers/common/codec/ya.make b/yql/essentials/providers/common/codec/ya.make index 13ccdc726b1..ca9d05dd2a2 100644 --- a/yql/essentials/providers/common/codec/ya.make +++ b/yql/essentials/providers/common/codec/ya.make @@ -30,6 +30,7 @@ END() RECURSE( arrow + yt_arrow_converter_interface ) RECURSE_FOR_TESTS( diff --git a/yql/essentials/providers/common/codec/yt_arrow_converter_interface/ya.make b/yql/essentials/providers/common/codec/yt_arrow_converter_interface/ya.make new file mode 100644 index 00000000000..de05d309851 --- /dev/null +++ b/yql/essentials/providers/common/codec/yt_arrow_converter_interface/ya.make @@ -0,0 +1,10 @@ +LIBRARY() + +PEERDIR( + contrib/libs/apache/arrow + yql/essentials/public/udf + yql/essentials/utils + yql/essentials/public/udf/arrow +) + +END() diff --git a/yql/essentials/providers/common/codec/yt_arrow_converter_interface/yt_arrow_converter.h b/yql/essentials/providers/common/codec/yt_arrow_converter_interface/yt_arrow_converter.h new file mode 100644 index 00000000000..532b9cfa8d2 --- /dev/null +++ b/yql/essentials/providers/common/codec/yt_arrow_converter_interface/yt_arrow_converter.h @@ -0,0 +1,13 @@ +#pragma once + +#include + +namespace NYql { + +class IYtColumnConverter { +public: + virtual arrow::Datum Convert(std::shared_ptr block) = 0; + virtual ~IYtColumnConverter() = default; +}; + +} diff --git a/yql/essentials/providers/common/codec/yt_arrow_converter_interface/yt_arrow_converter_details.h b/yql/essentials/providers/common/codec/yt_arrow_converter_interface/yt_arrow_converter_details.h new file mode 100644 index 00000000000..5aabd096825 --- /dev/null +++ b/yql/essentials/providers/common/codec/yt_arrow_converter_interface/yt_arrow_converter_details.h @@ -0,0 +1,100 @@ +#pragma once + +#include + +#include + +#include +#include + +namespace NYql { + +using namespace NYson::NDetail; + +class TYsonBuffer { +public: + TYsonBuffer(const std::string_view& s) : Data_(s.data()), Available_(s.size()) {} + + constexpr char Next() { + YQL_ENSURE(Available_-- > 0); + return *(++Data_); + } + + constexpr char Current() { + return *Data_; + } + + ui32 ReadVarUI32() { + char prev = Current(); + if (Y_LIKELY(!(prev & 0x80))) { + Next(); + return prev; + } + + return ReadVarSlow(); + } + + ui64 ReadVarUI64() { + char prev = Current(); + if (Y_LIKELY(!(prev & 0x80))) { + Next(); + return prev; + } + + return ReadVarSlow(); + } + + i32 ReadVarI32() { + return NYson::ZigZagDecode32(ReadVarUI32()); + } + + i64 ReadVarI64() { + return NYson::ZigZagDecode64(ReadVarUI64()); + } + + double NextDouble() { + double val = *reinterpret_cast(Data_); + Data_ += sizeof(double); + return val; + } + + void Skip(i32 cnt) { + Data_ += cnt; + } + + const char* Data() { + return Data_; + } + + size_t Available() const { + return Available_; + } +private: + template + constexpr T ReadVarSlow() { + T shift = 0; + T value = Current() & 0x7f; + for (;;) { + shift += 7; + value |= T(Next() & 0x7f) << shift; + if (!(Current() & 0x80)) { + break; + } + } + Next(); + return value; + } + + const char* Data_; + size_t Available_; +}; + + +class IYsonComplexTypeReader { +public: + using TPtr = std::unique_ptr; + virtual ~IYsonComplexTypeReader() = default; + virtual NUdf::TBlockItem GetItem(TYsonBuffer& buf) = 0; +}; + +} diff --git a/yql/essentials/sql/pg_dummy/pg_sql_dummy.cpp b/yql/essentials/sql/pg_dummy/pg_sql_dummy.cpp index b7f5fd713ae..9a9a10bab63 100644 --- a/yql/essentials/sql/pg_dummy/pg_sql_dummy.cpp +++ b/yql/essentials/sql/pg_dummy/pg_sql_dummy.cpp @@ -551,4 +551,14 @@ ui64 HexEncode(const char *src, size_t len, char *dst) { throw yexception() << "HexEncode in pg_dummy does nothing"; } + +std::unique_ptr BuildPgTopLevelColumnReader(std::unique_ptr&& builder, const NKikimr::NMiniKQL::TPgType* targetType) { + throw yexception() << "PG types are not supported"; +} + + +std::unique_ptr BuildPgYsonColumnReader(const NUdf::TPgTypeDescription& desc) { + throw yexception() << "PG types are not supported"; +} + } // NYql diff --git a/yt/yql/providers/yt/codec/README.md b/yt/yql/providers/yt/codec/README.md new file mode 100644 index 00000000000..f9157aec184 --- /dev/null +++ b/yt/yql/providers/yt/codec/README.md @@ -0,0 +1,4 @@ +# Примеры значений для типов с разными вариациями для YT колоночных конвертеров + +Тип(-ы) | YQL-выражение | YT | YT (+useNativeTypes) +----------------------------------------------- diff --git a/yt/yql/providers/yt/codec/ya.make b/yt/yql/providers/yt/codec/ya.make index 0bbae8552ba..8853ac14852 100644 --- a/yt/yql/providers/yt/codec/ya.make +++ b/yt/yql/providers/yt/codec/ya.make @@ -23,6 +23,7 @@ PEERDIR( yql/essentials/minikql/computation yql/essentials/public/udf yql/essentials/utils + yql/essentials/parser/pg_wrapper/interface yql/essentials/public/result_format yql/essentials/public/udf/arrow yql/essentials/providers/common/codec @@ -32,6 +33,7 @@ PEERDIR( yt/yql/providers/yt/common yt/yql/providers/yt/lib/mkql_helpers yt/yql/providers/yt/lib/skiff + yql/essentials/providers/common/codec/yt_arrow_converter_interface ) IF (MKQL_DISABLE_CODEGEN) diff --git a/yt/yql/providers/yt/codec/yt_arrow_converter.cpp b/yt/yql/providers/yt/codec/yt_arrow_converter.cpp index ea25ba325d5..f9cd95019a3 100644 --- a/yt/yql/providers/yt/codec/yt_arrow_converter.cpp +++ b/yt/yql/providers/yt/codec/yt_arrow_converter.cpp @@ -1,13 +1,15 @@ #include "yt_arrow_converter.h" -#include "yt_arrow_converter_details.h" +#include #include #include #include #include +#include #include #include #include +#include #include #include @@ -36,6 +38,39 @@ struct TYtColumnConverterSettings { }; } +template +class IYsonYQLComplexTypeReader : public IYsonComplexTypeReader { +public: + virtual NUdf::TBlockItem GetItem(TYsonBuffer& buf) = 0; + virtual NUdf::TBlockItem GetNotNull(TYsonBuffer&) = 0; + NUdf::TBlockItem GetNullableItem(TYsonBuffer& buf) { + char prev = buf.Current(); + if constexpr (Native) { + if (prev == EntitySymbol) { + buf.Next(); + return NUdf::TBlockItem(); + } + return GetNotNull(buf).MakeOptional(); + } + buf.Next(); + if (prev == EntitySymbol) { + return NUdf::TBlockItem(); + } + YQL_ENSURE(prev == BeginListSymbol); + if (buf.Current() == EndListSymbol) { + buf.Next(); + return NUdf::TBlockItem(); + } + auto result = GetNotNull(buf); + if (buf.Current() == ListItemSeparatorSymbol) { + buf.Next(); + } + YQL_ENSURE(buf.Current() == EndListSymbol); + buf.Next(); + return result.MakeOptional(); + } +}; + std::string_view GetNotNullString(auto& data, i64 idx) { i32 len; auto ptr = reinterpret_cast(data.GetValue(idx, &len)); @@ -236,10 +271,9 @@ NUdf::TBlockItem ReadYson(TYsonBuffer& buf) { } template -class TTupleYsonReader final : public IYsonComplexTypeReader { +class TTupleYsonReader final : public IYsonYQLComplexTypeReader { public: - using TIReaderPtr = std::unique_ptr>; - TTupleYsonReader(TVector&& children) + TTupleYsonReader(TVector&& children) : Children_(std::move(children)) , Items_(Children_.size()) {} @@ -264,12 +298,12 @@ public: return NUdf::TBlockItem(Items_.data()); } private: - const TVector Children_; + const TVector Children_; TVector Items_; }; template -class TStringYsonReader final : public IYsonComplexTypeReader { +class TStringYsonReader final : public IYsonYQLComplexTypeReader { public: NUdf::TBlockItem GetItem(TYsonBuffer& buf) override final { if constexpr (Nullable) { @@ -293,7 +327,7 @@ public: }; template -class TTzDateYsonReader final : public IYsonComplexTypeReader { +class TTzDateYsonReader final : public IYsonYQLComplexTypeReader { public: NUdf::TBlockItem GetItem(TYsonBuffer& buf) override final { if constexpr (Nullable) { @@ -334,7 +368,7 @@ public: }; template -class TFixedSizeYsonReader final : public IYsonComplexTypeReader { +class TFixedSizeYsonReader final : public IYsonYQLComplexTypeReader { public: NUdf::TBlockItem GetItem(TYsonBuffer& buf) override final { if constexpr (Nullable) { @@ -380,10 +414,9 @@ public: }; template -class TExternalOptYsonReader final : public IYsonComplexTypeReader { +class TExternalOptYsonReader final : public IYsonYQLComplexTypeReader { public: - using TIReaderPtr = std::unique_ptr>; - TExternalOptYsonReader(TIReaderPtr&& inner) + TExternalOptYsonReader(IYsonComplexTypeReader::TPtr&& inner) : Underlying_(std::move(inner)) {} @@ -413,12 +446,12 @@ public: Y_ABORT("Can't be called"); } private: - TIReaderPtr Underlying_; + IYsonComplexTypeReader::TPtr Underlying_; }; template struct TComplexTypeYsonReaderTraits { - using TResult = IYsonComplexTypeReader; + using TResult = IYsonComplexTypeReader; template using TTuple = TTupleYsonReader; // TODO: Implement reader for decimals @@ -429,7 +462,7 @@ struct TComplexTypeYsonReaderTraits { using TExtOptional = TExternalOptYsonReader; static std::unique_ptr MakePg(const NUdf::TPgTypeDescription& desc, const NUdf::IPgBuilder* pgBuilder) { - ythrow yexception() << "Complex type Yson reader not implemented for block resources"; + return BuildPgYsonColumnReader(desc); } static std::unique_ptr MakeResource(bool) { @@ -458,16 +491,18 @@ Y_FORCE_INLINE void AddFromYson(auto& reader, auto& builder, std::string_view ys builder->Add(std::move(res)); } -template -class TComplexTypeYsonColumnConverter { +template +class TComplexTypeYsonColumnConverter final : public IYtColumnConverter { public: - TComplexTypeYsonColumnConverter(TYtColumnConverterSettings& settings) : Settings_(settings) { + TComplexTypeYsonColumnConverter(TYtColumnConverterSettings&& settings) : Settings_(std::move(settings)) { Reader_ = NUdf::MakeBlockReaderImpl>(TTypeInfoHelper(), settings.Type, settings.PgBuilder); } arrow::Datum Convert(std::shared_ptr block) { auto& builder = Settings_.Builder; - if constexpr(!IsDictionary) { + if (block->type->id() != arrow::Type::DICTIONARY) { + // complex types comes in yson (which is binary type) + YQL_ENSURE(block->type->id() == arrow::Type::BINARY); arrow::BinaryArray binary(block); if (block->GetNullCount()) { for (i64 i = 0; i < block->length; ++i) { @@ -484,6 +519,10 @@ public: } return builder->Build(false); } + + // complex types comes in yson (which is binary type) + auto blockType = static_cast(*block->type).value_type()->id(); + YQL_ENSURE(blockType == arrow::Type::BINARY); arrow::DictionaryArray dict(block); arrow::BinaryArray binary(block->dictionary); auto data = dict.indices()->data()->GetValues(1); @@ -505,74 +544,78 @@ public: private: std::shared_ptr::TResult> Reader_; - TYtColumnConverterSettings& Settings_; + TYtColumnConverterSettings Settings_; }; -template -class TYtColumnConverter final : public IYtColumnConverter { +class TTopLevelYsonYtConverter final : public IYtColumnConverter { public: - TYtColumnConverter(TYtColumnConverterSettings&& settings) + TTopLevelYsonYtConverter(TYtColumnConverterSettings&& settings) : Settings_(std::move(settings)) - , DictYsonConverter_(Settings_) - , YsonConverter_(Settings_) - , DictPrimitiveConverter_(Settings_) , TopLevelYsonDictConverter_(Settings_) , TopLevelYsonConverter_(Settings_) - { - auto type = Settings_.Type; - IsJson_ = type->IsData() && AS_TYPE(TDataType, type)->GetDataSlot() == NUdf::EDataSlot::Json - || (Native && type->IsOptional() && AS_TYPE(TOptionalType, type)->GetItemType()->IsData() - && AS_TYPE(TDataType, AS_TYPE(TOptionalType, type)->GetItemType())->GetDataSlot() == NUdf::EDataSlot::Json); + {} + + arrow::Datum Convert(std::shared_ptr block) override { + if (arrow::Type::DICTIONARY == block->type->id()) { + return TopLevelYsonDictConverter_.Convert(block); + } else { + return TopLevelYsonConverter_.Convert(block); + } } +private: + TYtColumnConverterSettings Settings_; + TPrimitiveColumnConverter TopLevelYsonDictConverter_; + TPrimitiveColumnConverter TopLevelYsonConverter_; +}; + +template +class TTopLevelSimpleCastConverter final : public IYtColumnConverter { +public: + TTopLevelSimpleCastConverter(TYtColumnConverterSettings&& settings) + : Settings_(std::move(settings)) + , DictPrimitiveConverter_(Settings_) + {} arrow::Datum Convert(std::shared_ptr block) override { if (arrow::Type::DICTIONARY == block->type->id()) { - auto valType = static_cast(*block->type).value_type(); - if (Settings_.IsTopLevelYson) { - return TopLevelYsonDictConverter_.Convert(block); - } else if (valType->Equals(Settings_.ArrowType)) { - // just unpack - return DictPrimitiveConverter_.Convert(block); - } else if (arrow::Type::UINT8 == Settings_.ArrowType->id() && arrow::Type::BOOL == valType->id()) { - // unpack an cast - auto result = arrow::compute::Cast(DictPrimitiveConverter_.Convert(block), Settings_.ArrowType); - YQL_ENSURE(result.ok()); - return *result; - } else if (IsJson_ && arrow::Type::STRING == Settings_.ArrowType->id() && arrow::Type::BINARY == valType->id()) { - auto result = arrow::compute::Cast(DictPrimitiveConverter_.Convert(block), Settings_.ArrowType); - YQL_ENSURE(result.ok()); - return *result; - } else { - return DictYsonConverter_.Convert(block); - } + auto blockType = static_cast(*block->type).value_type(); + YQL_ENSURE(Expected == blockType->id()); + auto result = arrow::compute::Cast(DictPrimitiveConverter_.Convert(block), Settings_.ArrowType); + YQL_ENSURE(result.ok()); + return *result; } else { auto blockType = block->type; - if (Settings_.IsTopLevelYson) { - return TopLevelYsonConverter_.Convert(block); - } else if (blockType->Equals(Settings_.ArrowType)) { - return block; - } else if (arrow::Type::UINT8 == Settings_.ArrowType->id() && arrow::Type::BOOL == blockType->id()) { - auto result = arrow::compute::Cast(arrow::Datum(*block), Settings_.ArrowType); - YQL_ENSURE(result.ok()); - return *result; - } else if (IsJson_ && arrow::Type::STRING == Settings_.ArrowType->id() && arrow::Type::BINARY == blockType->id()) { - auto result = arrow::compute::Cast(arrow::Datum(*block), Settings_.ArrowType); - YQL_ENSURE(result.ok()); - return *result; - } else { - YQL_ENSURE(arrow::Type::BINARY == blockType->id()); - return YsonConverter_.Convert(block); - } + YQL_ENSURE(Expected == blockType->id()); + auto result = arrow::compute::Cast(arrow::Datum(*block), Settings_.ArrowType); + YQL_ENSURE(result.ok()); + return *result; + } + } +private: + TYtColumnConverterSettings Settings_; + TPrimitiveColumnConverter DictPrimitiveConverter_; +}; + +class TTopLevelAsIsConverter final : public IYtColumnConverter { +public: + TTopLevelAsIsConverter(TYtColumnConverterSettings&& settings) + : Settings_(std::move(settings)) + , DictPrimitiveConverter_(Settings_) + {} + + arrow::Datum Convert(std::shared_ptr block) override { + if (arrow::Type::DICTIONARY == block->type->id()) { + auto blockType = static_cast(*block->type).value_type(); + YQL_ENSURE(blockType->Equals(Settings_.ArrowType)); + return DictPrimitiveConverter_.Convert(block); + } else { + YQL_ENSURE(block->type->Equals(Settings_.ArrowType)); + return block; } } private: TYtColumnConverterSettings Settings_; - TComplexTypeYsonColumnConverter DictYsonConverter_; - TComplexTypeYsonColumnConverter YsonConverter_; TPrimitiveColumnConverter DictPrimitiveConverter_; - TPrimitiveColumnConverter TopLevelYsonDictConverter_; - TPrimitiveColumnConverter TopLevelYsonConverter_; - bool IsJson_; }; TYtColumnConverterSettings::TYtColumnConverterSettings(TType* type, const NUdf::IPgBuilder* pgBuilder, arrow::MemoryPool& pool, bool isNative) @@ -601,22 +644,64 @@ TYtColumnConverterSettings::TYtColumnConverterSettings(TType* type, const NUdf:: template typename T, typename Args, bool... Acc> struct TBoolDispatcher { + std::unique_ptr Dispatch(Args&& args) const { + return std::make_unique>(std::forward(args)); + } - std::unique_ptr Dispatch(Args&& args) const { - return std::make_unique>(std::forward(args)); - } - - template - auto Dispatch(Args&& args, bool head, Bools... tail) const { - return head ? - TBoolDispatcher().Dispatch(std::forward(args), tail...) : - TBoolDispatcher().Dispatch(std::forward(args), tail...); - } + template + auto Dispatch(Args&& args, bool head, Bools... tail) const { + return head ? + TBoolDispatcher().Dispatch(std::forward(args), tail...) : + TBoolDispatcher().Dispatch(std::forward(args), tail...); + } }; -std::unique_ptr MakeYtColumnConverter(NKikimr::NMiniKQL::TType* type, const NUdf::IPgBuilder* pgBuilder, arrow::MemoryPool& pool, bool isNative) { +std::unique_ptr MakeYtColumnConverter(TType* type, const NUdf::IPgBuilder* pgBuilder, arrow::MemoryPool& pool, bool isNative) { TYtColumnConverterSettings settings(type, pgBuilder, pool, isNative); bool isTopOptional = settings.IsTopOptional; - return TBoolDispatcher().Dispatch(std::move(settings), isNative, isTopOptional); + auto requestedType = type; + if (isTopOptional) { + requestedType = static_cast(type)->GetItemType(); + } + + if (type->IsPg()) { + // top-level pg is T? where T is int/string + return BuildPgTopLevelColumnReader(std::move(settings.Builder), static_cast(type)); + } + + if (type->IsData() && static_cast(type)->GetDataSlot() == NUdf::EDataSlot::Yson) { + // Special case: YT now has no non-optional Yson support + return std::make_unique(std::move(settings)); + } + + if (requestedType->IsData()) { + // T, T? where T is data + // There is no difference in native/non-native optional/non-optional + switch (*static_cast(requestedType)->GetDataSlot()) { + case NUdf::EDataSlot::Bool: + // YT type for bool is arrow::Type::BOOL, but yql type is arrow::Type::UINT8 + return std::make_unique>(std::move(settings)); + case NUdf::EDataSlot::String: + case NUdf::EDataSlot::Json: + case NUdf::EDataSlot::Yson: // Yson there is top-level optional + // YT type for Yson, Json, String is arrow::Type::BINARY, but yql type is arrow::Type::String + return std::make_unique>(std::move(settings)); + case NUdf::EDataSlot::Double: + case NUdf::EDataSlot::Int8: + case NUdf::EDataSlot::Uint8: + case NUdf::EDataSlot::Int16: + case NUdf::EDataSlot::Uint16: + case NUdf::EDataSlot::Int32: + case NUdf::EDataSlot::Uint32: + case NUdf::EDataSlot::Int64: + case NUdf::EDataSlot::Uint64: + // As is, except dictionary has come (in that way just unpack it) + return std::make_unique(std::move(settings)); + default: + Y_ABORT("That dataslot isn't supported (or implemented yet)"); + } + } + // Complex type and/or 2+ optional levels + return TBoolDispatcher().Dispatch(std::move(settings), isNative, isTopOptional); } } diff --git a/yt/yql/providers/yt/codec/yt_arrow_converter.h b/yt/yql/providers/yt/codec/yt_arrow_converter.h index a6f5c5a52e7..e54393131ac 100644 --- a/yt/yql/providers/yt/codec/yt_arrow_converter.h +++ b/yt/yql/providers/yt/codec/yt_arrow_converter.h @@ -1,22 +1,11 @@ +#pragma once +#include #include -#include -#include -#include -#include -#include -#include -#include -#include -#include +namespace NKikimr::NMiniKQL { +class TType; +} namespace NYql { - -class IYtColumnConverter { -public: - virtual arrow::Datum Convert(std::shared_ptr block) = 0; - virtual ~IYtColumnConverter() = default; -}; - std::unique_ptr MakeYtColumnConverter(NKikimr::NMiniKQL::TType* type, const NUdf::IPgBuilder* pgBuilder, arrow::MemoryPool& pool, bool isNative); } diff --git a/yt/yql/providers/yt/codec/yt_arrow_converter_details.h b/yt/yql/providers/yt/codec/yt_arrow_converter_details.h deleted file mode 100644 index ee59ab6b88f..00000000000 --- a/yt/yql/providers/yt/codec/yt_arrow_converter_details.h +++ /dev/null @@ -1,129 +0,0 @@ -#pragma once - -#include - -#include - -#include -#include - - -namespace NYql { - -using namespace NYson::NDetail; - -class TYsonBuffer { -public: - TYsonBuffer(const std::string_view& s) : Data_(s.data()), Available_(s.size()) {} - - constexpr char Next() { - YQL_ENSURE(Available_-- > 0); - return *(++Data_); - } - - constexpr char Current() { - return *Data_; - } - - ui32 ReadVarUI32() { - char prev = Current(); - if (Y_LIKELY(!(prev & 0x80))) { - Next(); - return prev; - } - - return ReadVarSlow(); - } - - ui64 ReadVarUI64() { - char prev = Current(); - if (Y_LIKELY(!(prev & 0x80))) { - Next(); - return prev; - } - - return ReadVarSlow(); - } - - i32 ReadVarI32() { - return NYson::ZigZagDecode32(ReadVarUI32()); - } - - i64 ReadVarI64() { - return NYson::ZigZagDecode64(ReadVarUI64()); - } - - double NextDouble() { - double val = *reinterpret_cast(Data_); - Data_ += sizeof(double); - return val; - } - - void Skip(i32 cnt) { - Data_ += cnt; - } - - const char* Data() { - return Data_; - } - - size_t Available() const { - return Available_; - } -private: - template - constexpr T ReadVarSlow() { - T shift = 0; - T value = Current() & 0x7f; - for (;;) { - shift += 7; - value |= T(Next() & 0x7f) << shift; - if (!(Current() & 0x80)) { - break; - } - } - Next(); - return value; - } - - const char* Data_; - size_t Available_; -}; - - -template -class IYsonComplexTypeReader { -public: - using TIReaderPtr = std::unique_ptr>; - virtual ~IYsonComplexTypeReader() = default; - virtual NUdf::TBlockItem GetItem(TYsonBuffer& buf) = 0; - virtual NUdf::TBlockItem GetNotNull(TYsonBuffer&) = 0; - NUdf::TBlockItem GetNullableItem(TYsonBuffer& buf) { - char prev = buf.Current(); - if constexpr (Native) { - if (prev == EntitySymbol) { - buf.Next(); - return NUdf::TBlockItem(); - } - return GetNotNull(buf).MakeOptional(); - } - buf.Next(); - if (prev == EntitySymbol) { - return NUdf::TBlockItem(); - } - YQL_ENSURE(prev == BeginListSymbol); - if (buf.Current() == EndListSymbol) { - buf.Next(); - return NUdf::TBlockItem(); - } - auto result = GetNotNull(buf); - if (buf.Current() == ListItemSeparatorSymbol) { - buf.Next(); - } - YQL_ENSURE(buf.Current() == EndListSymbol); - buf.Next(); - return result.MakeOptional(); - } -}; - -} diff --git a/yt/yql/providers/yt/codec/yt_codec_io.cpp b/yt/yql/providers/yt/codec/yt_codec_io.cpp index 598c1909df2..0530bba5f72 100644 --- a/yt/yql/providers/yt/codec/yt_codec_io.cpp +++ b/yt/yql/providers/yt/codec/yt_codec_io.cpp @@ -43,6 +43,8 @@ #include #include +#include +#include #include #include diff --git a/yt/yql/providers/yt/comp_nodes/dq/dq_yt_block_reader.cpp b/yt/yql/providers/yt/comp_nodes/dq/dq_yt_block_reader.cpp index 8b3019ffa37..b06f9034756 100644 --- a/yt/yql/providers/yt/comp_nodes/dq/dq_yt_block_reader.cpp +++ b/yt/yql/providers/yt/comp_nodes/dq/dq_yt_block_reader.cpp @@ -309,6 +309,9 @@ public: } arrow::Status OnRecordBatchDecoded(std::shared_ptr batch) override { + NKikimr::NMiniKQL::TScopedAlloc scope(__LOCATION__); + TThrowingBindTerminator t; + YQL_ENSURE(batch); MKQL_ADD_STAT(JobStats_, BlockCount, 1); std::vector result; -- cgit v1.3