YT-20113: Use code generated comparator in SimpleSort, PartitionSort if possible.

e23ff70d06989bb2b299bd02b1ac3c6c8f09b4e4

5 files changed, 45 insertions, 14 deletions

diff --git a/yt/yt/client/table_client/comparator.cpp b/yt/yt/client/table_client/comparator.cpp
index 1d8375387b..f49bec0fba 100644
--- a/yt/yt/client/table_client/comparator.cpp
+++ b/yt/yt/client/table_client/comparator.cpp
@@ -19,8 +19,9 @@ static const TLogger Logger("TableClientComparator");
 
 ////////////////////////////////////////////////////////////////////////////////
 
-TComparator::TComparator(std::vector<ESortOrder> sortOrders)
+TComparator::TComparator(std::vector<ESortOrder> sortOrders, TCallback<TUUComparerSignature> CGComparator)
     : SortOrders_(std::move(sortOrders))
+    , CGComparator_(CGComparator)
 { }
 
 void TComparator::Persist(const TPersistenceContext& context)
@@ -249,6 +250,19 @@ int TComparator::CompareKeys(const TKey& lhs, const TKey& rhs) const
     ValidateKey(lhs);
     ValidateKey(rhs);
 
+    if (CGComparator_) {
+        // Compare keys with code generated comparator.
+        int comparisonResult = CGComparator_.Run(lhs.Begin(), rhs.Begin(), GetLength());
+
+        if (comparisonResult == 0) {
+            return comparisonResult;
+        }
+
+        int differentColumnIdx = abs(comparisonResult) - 1;
+        int orderSign = SortOrders_[differentColumnIdx] == ESortOrder::Ascending ? +1 : -1;
+        return comparisonResult * orderSign;
+    }
+
     for (int index = 0; index < lhs.GetLength(); ++index) {
         auto valueComparisonResult = CompareValues(index, lhs[index], rhs[index]);
         if (valueComparisonResult != 0) {
diff --git a/yt/yt/client/table_client/comparator.h b/yt/yt/client/table_client/comparator.h
index 9c9e1d78d8..b228356555 100644
--- a/yt/yt/client/table_client/comparator.h
+++ b/yt/yt/client/table_client/comparator.h
@@ -28,7 +28,7 @@ public:
 
 public:
     TComparator() = default;
-    explicit TComparator(std::vector<ESortOrder> sortOrders);
+    TComparator(std::vector<ESortOrder> sortOrders, TCallback<TUUComparerSignature> CGComparator = {});
 
     void Persist(const TPersistenceContext& context);
 
@@ -85,6 +85,10 @@ public:
     explicit operator bool() const;
 
 private:
+    // Compiler generated comparer that is used in CompareKeys().
+    TCallback<TUUComparerSignature> CGComparator_;
+
+private:
     void ValidateKey(const TKey& key) const;
     void ValidateKeyBound(const TKeyBound& keyBound) const;
 };
@@ -96,7 +100,7 @@ void Serialize(const TComparator& comparator, NYson::IYsonConsumer* consumer);
 
 ////////////////////////////////////////////////////////////////////////////////
 
-using TPrefixComparer = int(const TUnversionedValue*, const TUnversionedValue*, int);
+using TPrefixComparer = TUUComparerSignature;
 
 ////////////////////////////////////////////////////////////////////////////////
 
diff --git a/yt/yt/client/table_client/public.h b/yt/yt/client/table_client/public.h
index 73b6c8dca6..f5285e898f 100644
--- a/yt/yt/client/table_client/public.h
+++ b/yt/yt/client/table_client/public.h
@@ -429,6 +429,9 @@ using TDynamicTableKeyMask = ui64;
 
 static_assert(sizeof(TDynamicTableKeyMask) * 8 == MaxKeyColumnCountInDynamicTable);
 
+// Function that compares two TUnversionedValue values.
+using TUUComparerSignature = int(const TUnversionedValue*, const TUnversionedValue*, int);
+
 ////////////////////////////////////////////////////////////////////////////////
 
 } // namespace NYT::NTableClient
diff --git a/yt/yt/client/table_client/schema.cpp b/yt/yt/client/table_client/schema.cpp
index 26b11be183..241c8019f2 100644
--- a/yt/yt/client/table_client/schema.cpp
+++ b/yt/yt/client/table_client/schema.cpp
@@ -749,6 +749,14 @@ bool TTableSchema::IsEmpty() const
     return Columns().empty();
 }
 
+bool TTableSchema::IsCGCompatarorApplicable() const
+{
+    auto keyTypes = GetKeyColumnTypes();
+    return std::none_of(keyTypes.begin(), keyTypes.end(), [] (auto type) {
+        return type == EValueType::Any;
+    });
+}
+
 std::optional<int> TTableSchema::GetTtlColumnIndex() const
 {
     auto* column = FindColumn(TtlColumnName);
@@ -1297,18 +1305,19 @@ TTableSchemaPtr TTableSchema::ToModifiedSchema(ETableSchemaModification schemaMo
     }
 }
 
-TComparator TTableSchema::ToComparator() const
+TComparator TTableSchema::ToComparator(TCallback<TUUComparerSignature> CGComparator) const
 {
-    if (!ColumnInfo_) {
-        return TComparator(std::vector<ESortOrder>());
-    }
-    const auto& info = *ColumnInfo_;
-    std::vector<ESortOrder> sortOrders(KeyColumnCount_);
-    for (int index = 0; index < KeyColumnCount_; ++index) {
-        YT_VERIFY(info.Columns[index].SortOrder());
-        sortOrders[index] = *info.Columns[index].SortOrder();
+    std::vector<ESortOrder> sortOrders;
+    if (ColumnInfo_) {
+        const auto& info = *ColumnInfo_;
+        sortOrders.resize(KeyColumnCount_);
+        for (int index = 0; index < KeyColumnCount_; ++index) {
+            YT_VERIFY(info.Columns[index].SortOrder());
+            sortOrders[index] = *info.Columns[index].SortOrder();
+        }
     }
-    return TComparator(std::move(sortOrders));
+
+    return TComparator(std::move(sortOrders), std::move(CGComparator));
 }
 
 void TTableSchema::Save(TStreamSaveContext& context) const
diff --git a/yt/yt/client/table_client/schema.h b/yt/yt/client/table_client/schema.h
index c425f244a2..8ece5019e9 100644
--- a/yt/yt/client/table_client/schema.h
+++ b/yt/yt/client/table_client/schema.h
@@ -286,6 +286,7 @@ public:
     bool IsUniqueKeys() const;
     bool HasRenamedColumns() const;
     bool IsEmpty() const;
+    bool IsCGCompatarorApplicable() const;
 
     std::optional<int> GetTtlColumnIndex() const;
 
@@ -380,7 +381,7 @@ public:
 
     TTableSchemaPtr ToModifiedSchema(ETableSchemaModification schemaModification) const;
 
-    TComparator ToComparator() const;
+    TComparator ToComparator(TCallback<TUUComparerSignature> CGComparator = {}) const;
 
     TKeyColumnTypes GetKeyColumnTypes() const;