Fix simplification of unique constraint. (#2950)

Author: Tony-Romanov

ydb/library/yql/ast/yql_constraint.cpp

@@ -1232,20 +1232,32 @@ TUniqueConstraintNodeBase<Distinct>::DoGetSimplifiedForType(const TTypeAnnotatio
  for (bool setChanged = true; setChanged;) {
  setChanged = false;
  for (auto it = sets.begin(); sets.end() != it;) {
- if (it->size() != 1U || it->front().size() <= 1U)
- ++it;
- else {
+ if (!it->empty() && it->front().size() > 1U) {
+ TPartOfConstraintBase::TSetType prefixes;
+ prefixes.reserve(it->size());
+ for (const auto& path : *it) {
+ if (path.size() > 1U) {
+ prefixes.emplace_back(getPrefix(path));
+ }
+ }
+
  auto from = it++;
- const auto prefix = getPrefix(from->front());
- while (sets.cend() != it && it->size() == 1U && it->front().size() > 1U && prefix == getPrefix(it->front()))
- ++it;
+ if (prefixes.size() < from->size())
+ continue;
 
- if (ssize_t(GetElementsCount(TBase::GetSubTypeByPath(prefix, rowType))) == std::distance(from, it)) {
- *from++ = TPartOfConstraintBase::TSetType{std::move(prefix)};
+ while (sets.cend() != it && it->size() == prefixes.size() &&
+ std::all_of(it->cbegin(), it->cend(), [&](const TPartOfConstraintBase::TPathType& path) { return path.size() > 1U && prefixes.contains(getPrefix(path)); })) {
+ ++it;
+ }
+
+ if (std::all_of(prefixes.cbegin(), prefixes.cend(),
+ [width = std::distance(from, it), &rowType] (const TPartOfConstraintBase::TPathType& path) { return width == ssize_t(GetElementsCount(TBase::GetSubTypeByPath(path, rowType))); })) {
+ *from++ =std::move(prefixes);
  it = sets.erase(from, it);
  changed = setChanged = true;
  }
- }
+ } else
+ ++it;
  }
  }
  }

ydb/library/yql/core/ut/yql_expr_constraint_ut.cpp

@@ -1566,6 +1566,36 @@ Y_UNIT_TEST_SUITE(TYqlExprConstraints) {
  CheckConstraint<TDistinctConstraintNode>(exprRoot, "Skip", "Distinct((key))");
  }
 
+ Y_UNIT_TEST(PartitionsByKeysWithCondense1WithPairItemsTupleKeyAndDuplicateOut) {
+ const auto s = R"(
+(
+ (let mr_sink (DataSink 'yt (quote plato)))
+ (let list (AsList
+ (AsStruct '('key '((Just (String '4)) (Just (String '%)))) '('subkey (Just (String 'c))) '('value (Just (String 'x))))
+ (AsStruct '('key '((Just (String '1)) (Just (String '€)))) '('subkey (Just (String 'b))) '('value (Just (String 'y))))
+ (AsStruct '('key '((Just (String '4)) (Just (String '$)))) '('subkey (Just (String 'b))) '('value (Just (String 'z))))
+ ))
+ (let extractor (lambda '(row) '((Nth (Member row 'key) '0) (Nth (Member row 'key) '1))))
+ (let aggr (PartitionsByKeys list (lambda '(item) (Member item 'key)) (Void) (Void)
+ (lambda '(stream) (Map (Condense1 stream extractor
+ (lambda '(row state) (IsKeySwitch row state extractor (lambda '(item) item)))
+ (lambda '(row state) state)
+ )
+ (lambda '(item) (AsStruct '('one '(item)) '('two '(item))))
+ ))
+ ))
+ (let world (Write! world mr_sink (Key '('table (String 'Output))) (Skip aggr (Uint64 '1)) '('('mode 'renew))))
+ (let world (Commit! world mr_sink))
+ (return world)
+)
+ )";
+
+ TExprContext exprCtx;
+ const auto exprRoot = ParseAndAnnotate(s, exprCtx);
+ CheckConstraint<TUniqueConstraintNode>(exprRoot, "Skip", "Unique(({one,two}))");
+ CheckConstraint<TDistinctConstraintNode>(exprRoot, "Skip", "Distinct(({one,two}))");
+ }
+
  Y_UNIT_TEST(ShuffleByKeysInputUnique) {
  const auto s = R"(
 (