YQL-16325 TupleLiteralBuilder

3 files changed, 53 insertions, 0 deletions

diff --git a/ydb/library/yql/minikql/mkql_node_builder.cpp b/ydb/library/yql/minikql/mkql_node_builder.cpp
index 920940aa96f..967e5bf0f53 100644
--- a/ydb/library/yql/minikql/mkql_node_builder.cpp
+++ b/ydb/library/yql/minikql/mkql_node_builder.cpp
@@ -106,6 +106,30 @@ TBlockType::EShape GetResultShape(const TVector<TType*>& types) {
     return result;
 }
 
+TTupleLiteralBuilder::TTupleLiteralBuilder(const TTypeEnvironment& env) : Env(env)
+{}
+
+void TTupleLiteralBuilder::Reserve(ui32 size) {
+    Types.reserve(size);
+    Values.reserve(size);
+}
+
+TTupleLiteralBuilder& TTupleLiteralBuilder::Add(TRuntimeNode value) {
+    Types.push_back(value.GetRuntimeType());
+    Values.push_back(value);
+    return *this;
+}
+
+TTupleLiteral* TTupleLiteralBuilder::Build() {
+    const auto& type = TTupleType::Create(Types.size(), Types.data(), Env);
+    return TTupleLiteral::Create(Values.size(), Values.data(), type, Env);
+}
+
+void TTupleLiteralBuilder::Clear() {
+    Values.clear();
+    Types.clear();
+}
+
 TStructTypeBuilder::TStructTypeBuilder(const TTypeEnvironment& env)
     : Env(&env)
 {
diff --git a/ydb/library/yql/minikql/mkql_node_builder.h b/ydb/library/yql/minikql/mkql_node_builder.h
index b854d539e72..ae4ea9a8b83 100644
--- a/ydb/library/yql/minikql/mkql_node_builder.h
+++ b/ydb/library/yql/minikql/mkql_node_builder.h
@@ -25,6 +25,21 @@ TDataType* UnpackOptionalData(TType* type, bool& isOptional);
 
 TBlockType::EShape GetResultShape(const TVector<TType*>& types);
 
+class TTupleLiteralBuilder {
+public:
+    TTupleLiteralBuilder(const TTypeEnvironment& env);
+    TTupleLiteralBuilder(const TTupleLiteralBuilder&) = default;
+    TTupleLiteralBuilder& operator=(const TTupleLiteralBuilder&) = default;
+    void Reserve(ui32 size);
+    TTupleLiteralBuilder& Add(TRuntimeNode value);
+    TTupleLiteral* Build();
+    void Clear();
+private:
+    const TTypeEnvironment& Env;
+    TVector<TType*> Types;
+    TVector<TRuntimeNode> Values;
+};
+
 class TStructTypeBuilder {
 public:
     TStructTypeBuilder(const TTypeEnvironment& env);
diff --git a/ydb/library/yql/minikql/mkql_node_builder_ut.cpp b/ydb/library/yql/minikql/mkql_node_builder_ut.cpp
index fcf32448f40..3a8b229326c 100644
--- a/ydb/library/yql/minikql/mkql_node_builder_ut.cpp
+++ b/ydb/library/yql/minikql/mkql_node_builder_ut.cpp
@@ -16,6 +16,20 @@ Y_UNIT_TEST_SUITE(TMiniKQLNodeBuilderTest) {
         UNIT_ASSERT_EQUAL(d2->GetType()->GetSchemeType(), NUdf::TDataType<ui32>::Id);
     }
 
+    Y_UNIT_TEST(TestTupleBuilder) {
+        TScopedAlloc alloc(__LOCATION__);
+        TTypeEnvironment env(alloc);
+        auto builder = TTupleLiteralBuilder(env);
+        const auto tuple = builder
+            .Add(TRuntimeNode(BuildDataLiteral(NUdf::TUnboxedValuePod(true), NUdf::EDataSlot::Bool, env), true))
+            .Add(TRuntimeNode(BuildDataLiteral(NUdf::TUnboxedValuePod(132), NUdf::EDataSlot::Uint32, env), true))
+            .Build();
+        UNIT_ASSERT_EQUAL(tuple->GetType()->GetKind(), TType::EKind::Tuple);
+        UNIT_ASSERT_VALUES_EQUAL(tuple->GetValuesCount(), 2);
+        builder.Clear();
+        UNIT_ASSERT_VALUES_EQUAL(builder.Build()->GetValuesCount(), 0);
+    }
+
     Y_UNIT_TEST(TestStructTypeBuilder) {
         TScopedAlloc alloc(__LOCATION__);
         TTypeEnvironment env(alloc);