[] Fix ExtractMembers optimizer - do not select syntetic columns as a 'light' column

1 files changed, 19 insertions, 2 deletions

diff --git a/ydb/library/yql/providers/s3/provider/yql_s3_logical_opt.cpp b/ydb/library/yql/providers/s3/provider/yql_s3_logical_opt.cpp
index 4f79b8fa72..e41a804261 100644
--- a/ydb/library/yql/providers/s3/provider/yql_s3_logical_opt.cpp
+++ b/ydb/library/yql/providers/s3/provider/yql_s3_logical_opt.cpp
@@ -468,8 +468,25 @@ public:
         const TStructExprType* readRowType =
             dqSource.Input().Maybe<TS3ParseSettings>().Cast().RowType().Ref().GetTypeAnn()->Cast<TTypeExprType>()->GetType()->Cast<TStructExprType>();
 
-        if (outputRowType->GetSize() == 0 && readRowType->GetSize() != 0) {
-            auto item = GetLightColumn(*readRowType);
+        TVector<const TItemExprType*> readRowDataItems = readRowType->GetItems();
+        TVector<const TItemExprType*> outputRowDataItems = outputRowType->GetItems();
+
+        if (auto settings = dqSource.Input().Maybe<TS3ParseSettings>().Cast().Settings()) {
+            if (auto ps = GetSetting(settings.Cast().Ref(), "partitionedby")) {
+                THashSet<TStringBuf> cols;
+                for (size_t i = 1; i < ps->ChildrenSize(); ++i) {
+                    YQL_ENSURE(ps->Child(i)->IsAtom());
+                    cols.insert(ps->Child(i)->Content());
+                }
+                auto isPartitionedBy = [&](const auto& item) { return cols.contains(item->GetName()); };
+                EraseIf(readRowDataItems, isPartitionedBy);
+                EraseIf(outputRowDataItems, isPartitionedBy);
+            }
+        }
+
+        if (outputRowDataItems.size() == 0 && readRowDataItems.size() != 0) {
+            const TStructExprType* readRowDataType = ctx.MakeType<TStructExprType>(readRowDataItems);
+            auto item = GetLightColumn(*readRowDataType);
             YQL_ENSURE(item);
             readRowType = ctx.MakeType<TStructExprType>(TVector<const TItemExprType*>{item});
         } else {