Merge 53b4046 into 8f381d4

ydb-platform · May 26, 2024 · ba3d1ba · ba3d1ba
2 parents 8f381d4 + 53b4046
commit ba3d1ba
Show file tree

Hide file tree

Showing 2 changed files with 21 additions and 2 deletions.
diff --git a/ydb/library/yql/providers/s3/provider/yql_s3_dq_integration.cpp b/ydb/library/yql/providers/s3/provider/yql_s3_dq_integration.cpp
@@ -15,8 +15,10 @@
 #include <ydb/library/yql/providers/s3/range_helpers/file_tree_builder.h>
 #include <ydb/library/yql/providers/s3/range_helpers/path_list_reader.h>
 #include <ydb/library/yql/utils/log/log.h>
+#include <ydb/library/yql/core/yql_opt_utils.h>
 
 #include <library/cpp/json/writer/json_value.h>
+#include <library/cpp/yson/node/node_io.h>
 
 namespace NYql {
 
@@ -179,7 +181,7 @@ class TS3DqIntegration: public TDqIntegrationBase {
     TMaybe<TOptimizerStatistics> ReadStatistics(const TExprNode::TPtr& sourceWrap, TExprContext& ctx) override {
         Y_UNUSED(ctx);
         double size = 0;
-        double cols = 0;
+        int cols = 0;
         double rows = 0;
         if (const auto& maybeParseSettings = TMaybeNode<TS3ParseSettings>(sourceWrap->Child(0))) {
             const auto& parseSettings = maybeParseSettings.Cast();
@@ -195,12 +197,28 @@ class TS3DqIntegration: public TDqIntegrationBase {
                 }
             }
 
+            TVector<TString>* primaryKey = nullptr;
+            if (auto constraints = GetSetting(parseSettings.Settings().Ref(), "constraints"sv)) {
+                auto node = NYT::NodeFromYsonString(constraints->Child(1)->Content());
+                auto* primaryKeyNode = node.AsMap().FindPtr("primary_key");
+                if (primaryKeyNode) {
+                    TVector<TString> parsed;
+                    for (auto col : primaryKeyNode->AsList()) {
+                        parsed.push_back(col.AsString());
+                    }
+                    State_->PrimaryKeys.emplace_back(std::move(parsed));
+                    primaryKey = &State_->PrimaryKeys.back();
+                }
+            }
+
             if (parseSettings.RowType().Maybe<TCoStructType>()) {
                 cols = parseSettings.RowType().Ptr()->ChildrenSize();
             }
 
             rows = size / 1024; // magic estimate
-            return TOptimizerStatistics(BaseTable, rows, cols, size);
+            return primaryKey 
+                ? TOptimizerStatistics(BaseTable, rows, cols, size, size, *primaryKey)
+                : TOptimizerStatistics(BaseTable, rows, cols, size, size);
         } else {
             return Nothing();
         }

diff --git a/ydb/library/yql/providers/s3/provider/yql_s3_provider.h b/ydb/library/yql/providers/s3/provider/yql_s3_provider.h
@@ -29,6 +29,7 @@ struct TS3State : public TThrRefBase
     ISecuredServiceAccountCredentialsFactory::TPtr CredentialsFactory;
     IHTTPGateway::TPtr Gateway;
     ui32 ExecutorPoolId = 0;
+    std::list<TVector<TString>> PrimaryKeys;
 };
 
 TDataProviderInitializer GetS3DataProviderInitializer(IHTTPGateway::TPtr gateway, ISecuredServiceAccountCredentialsFactory::TPtr credentialsFactory = nullptr, bool allowLocalFiles = false);