YQL-19336: Fix aliased columns in PG

YQL-19336: Fix aliased columns in PG
commit_hash:30ac6e8d9d0876fed9ba1e8b848d5c7ad8826dda

5 files changed, 49 insertions, 21 deletions

diff --git a/yql/essentials/core/common_opt/yql_co_pgselect.cpp b/yql/essentials/core/common_opt/yql_co_pgselect.cpp
index 93b307aab7..e8d37492a8 100644
--- a/yql/essentials/core/common_opt/yql_co_pgselect.cpp
+++ b/yql/essentials/core/common_opt/yql_co_pgselect.cpp
@@ -1022,7 +1022,7 @@ TUsedColumns GatherUsedColumns(const TExprNode::TPtr& result, const TExprNode::T
 void FillInputIndices(const TExprNode::TPtr& from, const TExprNode::TPtr& finalExtTypes,
     TUsedColumns& usedColumns, TOptimizeContext& optCtx) {
     for (auto& x : usedColumns) {
-        TStringBuf alias;
+        TString alias;
         TStringBuf column = NTypeAnnImpl::RemoveAlias(x.first, alias);
 
         bool foundColumn = false;
diff --git a/yql/essentials/core/type_ann/type_ann_pg.cpp b/yql/essentials/core/type_ann/type_ann_pg.cpp
index d0d73f22bb..78a035b70e 100644
--- a/yql/essentials/core/type_ann/type_ann_pg.cpp
+++ b/yql/essentials/core/type_ann/type_ann_pg.cpp
@@ -346,7 +346,7 @@ IGraphTransformer::TStatus PgCallWrapper(const TExprNode::TPtr& input, TExprNode
             TVector<const TItemExprType*> items;
             for (size_t i = 0; i < proc.OutputArgTypes.size(); ++i) {
                 items.push_back(ctx.Expr.MakeType<TItemExprType>(
-                    resultColumnOrder->AddColumn(proc.OutputArgNames[i]), 
+                    resultColumnOrder->AddColumn(proc.OutputArgNames[i]),
                     ctx.Expr.MakeType<TPgExprType>(proc.OutputArgTypes[i])));
             }
 
@@ -697,7 +697,7 @@ IGraphTransformer::TStatus PgArrayOpWrapper(const TExprNode::TPtr& input, TExprN
 
         const auto& opResDesc = NPg::LookupType(oper.ResultType);
         if (opResDesc.Name != "bool") {
-            ctx.Expr.AddError(TIssue(ctx.Expr.GetPosition(input->Pos()), TStringBuilder() << 
+            ctx.Expr.AddError(TIssue(ctx.Expr.GetPosition(input->Pos()), TStringBuilder() <<
                 "Expected boolean operator result, but got: " << opResDesc.Name));
             return IGraphTransformer::TStatus::Error;
         }
@@ -1021,7 +1021,7 @@ IGraphTransformer::TStatus PgNullIfWrapper(const TExprNode::TPtr& input, TExprNo
         input->ChildRef(0) = WrapWithPgCast(std::move(input->ChildRef(0)), commonType->TypeId, ctx.Expr);
         return IGraphTransformer::TStatus::Repeat;
     }
-    
+
     input->SetTypeAnn(ctx.Expr.MakeType<TPgExprType>(commonType->TypeId));
     return IGraphTransformer::TStatus::Ok;
 }
@@ -1144,7 +1144,7 @@ IGraphTransformer::TStatus PgReplaceUnknownWrapper(const TExprNode::TPtr& input,
 
     auto structType = listType->GetItemType()->Cast<TStructExprType>();
     auto structItemTypes = structType->GetItems();
-    const bool noUnknowns = std::none_of(structItemTypes.cbegin(), structItemTypes.cend(), 
+    const bool noUnknowns = std::none_of(structItemTypes.cbegin(), structItemTypes.cend(),
         [] (const TItemExprType* item) {
                 const auto* itemType = item->GetItemType();
                 return itemType->GetKind() == ETypeAnnotationKind::Pg && itemType->Cast<TPgExprType>()->GetId() == NPg::UnknownOid;
@@ -2010,12 +2010,23 @@ bool ValidateWindowRefs(const TExprNode::TPtr& root, const TExprNode* windows, T
     return !isError;
 }
 
+TString EscapeDotsInAlias(TStringBuf alias) {
+    TStringBuilder sb;
+    for (auto c: alias) {
+        if (c == '.' || c == '\\') {
+            sb << '\\';
+        }
+        sb << c;
+    }
+    return sb;
+}
+
 TString MakeAliasedColumn(TStringBuf alias, TStringBuf column) {
     if (!alias) {
         return TString(column);
     }
 
-    return TStringBuilder() << "_alias_" << alias << "." << column;
+    return TStringBuilder() << "_alias_" << EscapeDotsInAlias(alias) << "." << column;
 }
 
 const TItemExprType* AddAlias(const TString& alias, const TItemExprType* item, TExprContext& ctx) {
@@ -2027,22 +2038,30 @@ const TItemExprType* AddAlias(const TString& alias, const TItemExprType* item, T
 }
 
 TStringBuf RemoveAlias(TStringBuf column) {
-    TStringBuf tmp;
+    TString tmp;
     return RemoveAlias(column, tmp);
 }
 
-TStringBuf RemoveAlias(TStringBuf column, TStringBuf& alias) {
+TStringBuf RemoveAlias(TStringBuf column, TString& alias) {
     if (!column.StartsWith("_alias_")) {
         alias = "";
         return column;
     }
-
-    auto columnPos = column.find('.', 7);
-    YQL_ENSURE(columnPos != TString::npos);
-    columnPos += 1;
-    YQL_ENSURE(columnPos != column.size());
-    alias = column.substr(7, columnPos - 7 - 1);
-    return column.substr(columnPos);
+    column = column.substr(7);
+    TStringBuilder aliasBuilder;
+    for (size_t i = 0; i < column.size(); ++i) {
+        if (column[i] == '\\') {
+            YQL_ENSURE(i + 1 < column.size());
+            aliasBuilder << column[++i];
+            continue;
+        }
+        if (column[i] == '.') {
+            alias = aliasBuilder;
+            return column.substr(i + 1);
+        }
+        aliasBuilder << column[i];
+    }
+    YQL_ENSURE(false, "No dot ('.') found in alised column");
 }
 
 const TItemExprType* RemoveAlias(const TItemExprType* item, TExprContext& ctx) {
@@ -2877,7 +2896,7 @@ bool ReplaceProjectionRefs(TExprNode::TPtr& lambda, const TStringBuf& scope, con
             for (ui32 i = 0; i < projectionOrders.size(); ++i) {
                 if (index >= current && index < current + projectionOrders[i]->first.Size()) {
                     TStringBuf column = projectionOrders[i]->first[index - current].PhysicalName;
-                    TStringBuf alias;
+                    TString alias;
                     column = RemoveAlias(column, alias);
 
                     if (result && projectionOrders[i]->second) {
@@ -3346,7 +3365,7 @@ bool GatherExtraSortColumns(const TExprNode& data, const TInputs& inputs, TExprN
                 }
 
                 if (node->IsCallable("Member") && &node->Head() == arg) {
-                    TStringBuf alias;
+                    TString alias;
                     TStringBuf column = NTypeAnnImpl::RemoveAlias(node->Tail().Content(), alias);
 
                     TMaybe<ui32> index;
@@ -3495,7 +3514,7 @@ IGraphTransformer::TStatus PgSetItemWrapper(const TExprNode::TPtr& input, TExprN
                             "Incorrect fill_target_columns option"));
                         return IGraphTransformer::TStatus::Error;
                     }
-		        }
+                }
                 else if (optionName == "unknowns_allowed") {
                     hasUnknownsAllowed = true;
                 }
@@ -4456,7 +4475,7 @@ IGraphTransformer::TStatus PgSetItemWrapper(const TExprNode::TPtr& input, TExprN
                                                 return IGraphTransformer::TStatus::Error;
                                             }
                                         }
-                                        
+
                                         if (rightSideUsing.contains(lcase)) {
                                             lrNames[1] = ctx.Expr.NewList(inp->Pos(), {ctx.Expr.NewAtom(inp->Pos(), rightSideUsing[lcase])});
                                         } else {
@@ -5442,7 +5461,7 @@ IGraphTransformer::TStatus PgArrayWrapper(const TExprNode::TPtr& input, TExprNod
 
     bool castsNeeded = false;
     const NPg::TTypeDesc* elemTypeDesc;
-    if (const auto issue = NPg::LookupCommonType(argTypes, 
+    if (const auto issue = NPg::LookupCommonType(argTypes,
         [&input, &ctx](size_t i) {
             return ctx.Expr.GetPosition(input->Child(i)->Pos());
         }, elemTypeDesc, castsNeeded))
diff --git a/yql/essentials/core/type_ann/type_ann_pg.h b/yql/essentials/core/type_ann/type_ann_pg.h
index 85ca8285ef..042d2c0092 100644
--- a/yql/essentials/core/type_ann/type_ann_pg.h
+++ b/yql/essentials/core/type_ann/type_ann_pg.h
@@ -12,7 +12,7 @@ TExprNodePtr WrapWithPgCast(TExprNodePtr&& node, ui32 targetTypeId, TExprContext
 TString MakeAliasedColumn(TStringBuf alias, TStringBuf column);
 const TItemExprType* AddAlias(const TString& alias, const TItemExprType* item, TExprContext& ctx);
 TStringBuf RemoveAlias(TStringBuf column);
-TStringBuf RemoveAlias(TStringBuf column, TStringBuf& alias);
+TStringBuf RemoveAlias(TStringBuf column, TString& alias);
 const TItemExprType* RemoveAlias(const TItemExprType* item, TExprContext& ctx);
 TMap<TString, ui32> ExtractExternalColumns(const TExprNode& select);
 bool IsPlainMemberOverArg(const TExprNode& expr, TStringBuf& memberName);
diff --git a/yql/essentials/tests/sql/sql2yql/canondata/result.json b/yql/essentials/tests/sql/sql2yql/canondata/result.json
index 5e6388be95..e73aa91e9e 100644
--- a/yql/essentials/tests/sql/sql2yql/canondata/result.json
+++ b/yql/essentials/tests/sql/sql2yql/canondata/result.json
@@ -12368,6 +12368,13 @@
             "uri": "https://{canondata_backend}/1937429/434276f26b2857be3c5ad3fdbbf877d2bf775ac5/resource.tar.gz#test_sql2yql.test_pg-aggregate_scalar_minus_zero_/sql.yql"
         }
     ],
+    "test_sql2yql.test[pg-aliased_columns]": [
+        {
+            "checksum": "7ca1f0f6ce1395c4d6f826660886df83",
+            "size": 953,
+            "uri": "https://{canondata_backend}/1777230/27c189b00ecc5d4153899123da0bf3c72d8cfd80/resource.tar.gz#test_sql2yql.test_pg-aliased_columns_/sql.yql"
+        }
+    ],
     "test_sql2yql.test[pg-all_data]": [
         {
             "checksum": "e2e983a696817fe3cbebcbbf79264fd1",
diff --git a/yql/essentials/tests/sql/suites/pg/aliased_columns.sql b/yql/essentials/tests/sql/suites/pg/aliased_columns.sql
new file mode 100644
index 0000000000..6ba1335386
--- /dev/null
+++ b/yql/essentials/tests/sql/suites/pg/aliased_columns.sql
@@ -0,0 +1,2 @@
+--!syntax_pg
+select a, b, c from (select 1 as a, 2 as b, 3 as c) as "A.B" order by b