YQL-16727 initial support of materialized set_of functions

Изменения в ydb/library/yql/sql/pg/ut PgOptimizer - веса зависят от work_mem

6 files changed, 149 insertions, 11 deletions

diff --git a/ydb/library/yql/parser/pg_wrapper/comp_factory.cpp b/ydb/library/yql/parser/pg_wrapper/comp_factory.cpp
index 2863050f4e7..6337bb8cf72 100644
--- a/ydb/library/yql/parser/pg_wrapper/comp_factory.cpp
+++ b/ydb/library/yql/parser/pg_wrapper/comp_factory.cpp
@@ -296,6 +296,14 @@ public:
     }
 
     ~TReturnSetInfo() {
+        if (Ref().expectedDesc) {
+            FreeTupleDesc(Ref().expectedDesc);
+        }
+
+        if (Ref().setResult) {
+            tuplestore_end(Ref().setResult);
+        }
+
         TWithDefaultMiniKQLAlloc::FreeWithSize(Ptr, sizeof(ReturnSetInfo));
     }
 
@@ -483,7 +491,15 @@ private:
             {
                 auto& callInfo = CallInfo.Ref();
                 callInfo.resultinfo = (fmNodePtr)&RSInfo.Ref();
-                ((ReturnSetInfo*)callInfo.resultinfo)->econtext = &ExprContextHolder.Ref();
+                auto& rsInfo = *(ReturnSetInfo*)callInfo.resultinfo;
+                rsInfo.econtext = &ExprContextHolder.Ref();
+                rsInfo.allowedModes = (int) (SFRM_ValuePerCall | SFRM_Materialize);
+                rsInfo.returnMode = SFRM_ValuePerCall;
+                rsInfo.setResult = nullptr;
+                rsInfo.setDesc = nullptr;
+                rsInfo.expectedDesc = CreateTemplateTupleDesc(1);
+                TupleDescInitEntry(rsInfo.expectedDesc, (AttrNumber) 1, nullptr, RetTypeDesc.TypeId, -1, 0);
+                TupleSlot = MakeSingleTupleTableSlot(rsInfo.expectedDesc, &TTSOpsMinimalTuple);
                 for (ui32 i = 0; i < args.size(); ++i) {
                     const auto& value = args[i];
                     NullableDatum argDatum = { 0, false };
@@ -500,6 +516,9 @@ private:
             }
 
             ~TIterator() {
+                if (TupleSlot) {
+                    ExecDropSingleTupleTableSlot(TupleSlot);
+                }
             }
 
         private:
@@ -508,19 +527,58 @@ private:
                     return false;
                 }
 
+                if (RSInfo.Ref().setResult) {
+                    return CopyTuple(value);
+                }
+
                 auto& callInfo = CallInfo.Ref();
                 callInfo.isnull = false;
                 auto ret = callInfo.flinfo->fn_addr(&callInfo);
-                if (RSInfo.Ref().isDone == ExprEndResult) {
+                if (RSInfo.Ref().returnMode == SFRM_Materialize) {
+                    Y_ENSURE(RSInfo.Ref().isDone == ExprSingleResult);
+                    Y_ENSURE(RSInfo.Ref().setResult);
+                    auto readPtr = tuplestore_alloc_read_pointer(RSInfo.Ref().setResult, 0);
+                    tuplestore_select_read_pointer(RSInfo.Ref().setResult, readPtr);
+                    return CopyTuple(value);
+                } else {
+                    if (RSInfo.Ref().isDone == ExprEndResult) {
+                        IsFinished = true;
+                        return false;
+                    }
+
+                    if (callInfo.isnull) {
+                        value = NUdf::TUnboxedValuePod();
+                    } else {
+                        value = AnyDatumToPod(ret, RetTypeDesc.PassByValue);
+                    }
+
+                    return true;
+                }
+            }
+
+            bool CopyTuple(NUdf::TUnboxedValue& value) {
+                if (!tuplestore_gettupleslot(RSInfo.Ref().setResult, true, false, TupleSlot)) {
                     IsFinished = true;
                     return false;
                 }
-
-                if (callInfo.isnull) {
+                
+                slot_getallattrs(TupleSlot);
+                Y_ENSURE(TupleSlot->tts_nvalid == 1);
+                if (TupleSlot->tts_isnull[0]) {
                     value = NUdf::TUnboxedValuePod();
                 } else {
-                    value = AnyDatumToPod(ret, RetTypeDesc.PassByValue);
-                }
+                    auto datum = TupleSlot->tts_values[0];
+                    if (RetTypeDesc.PassByValue) {
+                        value = ScalarDatumToPod(datum);
+                    } else if (RetTypeDesc.TypeLen == -1) {
+                        const text* orig = (const text*)datum;
+                        value = PointerDatumToPod((Datum)MakeVar(GetVarBuf(orig)));
+                    } else {
+                        Y_ENSURE(RetTypeDesc.TypeLen == -2);
+                        const char* orig = (const char*)datum;
+                        value = PointerDatumToPod((Datum)MakeCString(orig));
+                    }
+               }
 
                 return true;
             }
@@ -533,6 +591,7 @@ private:
             TFunctionCallInfo CallInfo;
             TReturnSetInfo RSInfo;
             bool IsFinished = false;
+            TupleTableSlot* TupleSlot = nullptr;
         };
 
         TListValue(TMemoryUsageInfo* memInfo, TComputationContext& compCtx,
diff --git a/ydb/library/yql/parser/pg_wrapper/functions.md b/ydb/library/yql/parser/pg_wrapper/functions.md
index 2afc792bf11..cfe1b2ce0c0 100644
--- a/ydb/library/yql/parser/pg_wrapper/functions.md
+++ b/ydb/library/yql/parser/pg_wrapper/functions.md
@@ -3078,9 +3078,9 @@ Table 9.47. JSON Processing Functions
 ||Function|Description|Example(s)||

 ||json_array_elements ( json ) → setof json  

 jsonb_array_elements ( jsonb ) → setof jsonb|

-Expands the top-level JSON array into a set of JSON values. (NOT SUPPORTED)|

+Expands the top-level JSON array into a set of JSON values.|

 ```sql

-#SELECT * FROM json_array_elements('[1,true, [2,false]]') as a → [

+SELECT * FROM json_array_elements('[1,true, [2,false]]') as a → [

 1

 true

 [2,false]

@@ -3088,9 +3088,9 @@ true
 ```||

 ||json_array_elements_text ( json ) → setof text  

 jsonb_array_elements_text ( jsonb ) → setof text|

-Expands the top-level JSON array into a set of text values. (NOT SUPPORTED)|

+Expands the top-level JSON array into a set of text values.|

 ```sql

-#SELECT * FROM json_array_elements_text('["foo", "bar"]') as a → [

+SELECT * FROM json_array_elements_text('["foo", "bar"]') as a → [

 foo

 bar

 ]

diff --git a/ydb/library/yql/parser/pg_wrapper/parser.cpp b/ydb/library/yql/parser/pg_wrapper/parser.cpp
index a22a0cc81ad..fc74e000ec7 100644
--- a/ydb/library/yql/parser/pg_wrapper/parser.cpp
+++ b/ydb/library/yql/parser/pg_wrapper/parser.cpp
@@ -16,12 +16,14 @@
 #undef SIZEOF_SIZE_T
 extern "C" {
 #include "postgres.h"
+#include "access/session.h"
 #include "access/xact.h"
 #include "mb/pg_wchar.h"
 #include "nodes/pg_list.h"
 #include "nodes/parsenodes.h"
 #include "nodes/value.h"
 #include "parser/parser.h"
+#include "utils/guc.h"
 #include "utils/palloc.h"
 #include "utils/memutils.h"
 #include "utils/memdebug.h"
@@ -257,4 +259,7 @@ extern "C" void setup_pg_thread_cleanup() {
     MyProc = (PGPROC*)malloc(sizeof(PGPROC));
     Zero(*MyProc);
     StartTransactionCommand();
+
+    InitializeSession();
+    work_mem = MAX_KILOBYTES; // a way to postpone spilling for tuple stores
 };
diff --git a/ydb/library/yql/parser/pg_wrapper/postgresql/src/backend/access/common/tupdesc.c b/ydb/library/yql/parser/pg_wrapper/postgresql/src/backend/access/common/tupdesc.c
index 4c63bd4dc64..a496d5e54f2 100644
--- a/ydb/library/yql/parser/pg_wrapper/postgresql/src/backend/access/common/tupdesc.c
+++ b/ydb/library/yql/parser/pg_wrapper/postgresql/src/backend/access/common/tupdesc.c
@@ -580,7 +580,7 @@ hashTupleDesc(TupleDesc desc)
  * TupleDescInitEntryCollation.
  */
 void
-TupleDescInitEntry(TupleDesc desc,
+TupleDescInitEntry_original(TupleDesc desc,
 				   AttrNumber attributeNumber,
 				   const char *attributeName,
 				   Oid oidtypeid,
diff --git a/ydb/library/yql/parser/pg_wrapper/type_cache.cpp b/ydb/library/yql/parser/pg_wrapper/type_cache.cpp
index 022df305723..a581cc8d132 100644
--- a/ydb/library/yql/parser/pg_wrapper/type_cache.cpp
+++ b/ydb/library/yql/parser/pg_wrapper/type_cache.cpp
@@ -9,6 +9,9 @@ extern "C" {
 #include "postgres.h"
 #include "catalog/pg_operator_d.h"
 #include "access/hash.h"
+#include "access/toast_compression.h"
+#include "access/tupdesc.h"
+#include "utils/builtins.h"
 #include "utils/catcache.h"
 #include "utils/fmgroids.h"
 #include "utils/typcache.h"
@@ -502,5 +505,68 @@ void destroy_typecache_hashtable() {
     YQL_TypeCacheHash = nullptr;
 }
 
+void
+TupleDescInitEntry(TupleDesc desc,
+                   AttrNumber attributeNumber,
+                   const char *attributeName,
+                   Oid oidtypeid,
+                   int32 typmod,
+                   int attdim)
+{
+    HeapTuple tuple;
+    Form_pg_attribute att;
+
+    /*
+     * sanity checks
+     */
+    AssertArg(PointerIsValid(desc));
+    AssertArg(attributeNumber >= 1);
+    AssertArg(attributeNumber <= desc->natts);
+
+    /*
+     * initialize the attribute fields
+     */
+    att = TupleDescAttr(desc, attributeNumber - 1);
+
+    att->attrelid = 0;			/* dummy value */
+
+    /*
+     * Note: attributeName can be NULL, because the planner doesn't always
+     * fill in valid resname values in targetlists, particularly for resjunk
+     * attributes. Also, do nothing if caller wants to re-use the old attname.
+     */
+    if (attributeName == NULL)
+        MemSet(NameStr(att->attname), 0, NAMEDATALEN);
+    else if (attributeName != NameStr(att->attname))
+        namestrcpy(&(att->attname), attributeName);
+
+    att->attstattarget = -1;
+    att->attcacheoff = -1;
+    att->atttypmod = typmod;
+
+    att->attnum = attributeNumber;
+    att->attndims = attdim;
+
+    att->attnotnull = false;
+    att->atthasdef = false;
+    att->atthasmissing = false;
+    att->attidentity = '\0';
+    att->attgenerated = '\0';
+    att->attisdropped = false;
+    att->attislocal = true;
+    att->attinhcount = 0;
+    /* attacl, attoptions and attfdwoptions are not present in tupledescs */
+
+    const auto& typeDesc = NYql::NPg::LookupType(oidtypeid);
+    att->atttypid = oidtypeid;
+    att->attlen = typeDesc.TypeLen;
+    att->attbyval = typeDesc.PassByValue;
+    att->attalign = typeDesc.TypeAlign;
+    att->attstorage = TYPSTORAGE_PLAIN;
+    att->attcompression = InvalidCompressionMethod;
+    att->attcollation = typeDesc.TypeCollation;
+}
+
+
 }
 }
\ No newline at end of file
diff --git a/ydb/library/yql/sql/pg/optimizer.cpp b/ydb/library/yql/sql/pg/optimizer.cpp
index 33f04365171..62e7019d75e 100644
--- a/ydb/library/yql/sql/pg/optimizer.cpp
+++ b/ydb/library/yql/sql/pg/optimizer.cpp
@@ -6,6 +6,7 @@
 #include <ydb/library/yql/utils/yql_panic.h>
 
 #include <util/string/builder.h>
+#include <util/generic/scope.h>
 
 #ifdef _WIN32
 #define __restrict
@@ -17,6 +18,7 @@
 
 extern "C" {
 #include "postgres.h"
+#include "miscadmin.h"
 #include "optimizer/paths.h"
 #include "nodes/print.h"
 #include "utils/selfuncs.h"
@@ -120,6 +122,12 @@ TPgOptimizer::~TPgOptimizer()
 TPgOptimizer::TOutput TPgOptimizer::JoinSearch()
 {
     TArenaMemoryContext ctx;
+    auto prev_work_mem = work_mem;
+    work_mem = 4096;
+    Y_DEFER {
+        work_mem = prev_work_mem;
+    };
+
     auto* rel = JoinSearchInternal();
     return MakeOutput(rel->cheapest_total_path);
 }