Merge ddc79c0 into 665c03b

Author: ivanmorozov333

ydb/core/formats/arrow/replace_key.h

@@ -57,6 +57,9 @@ class TReplaceKeyTemplate {
         , Position(position)
     {
         Y_ABORT_UNLESS(Size() > 0 && Position < (ui64)Column(0).length());
+        for (auto&& i : *Columns) {
+            Types.emplace_back(i->type_id());
+        }
     }
 
     template<typename T = TArrayVecPtr> requires IsOwning
@@ -65,13 +68,16 @@ class TReplaceKeyTemplate {
         , Position(position)
     {
         Y_ABORT_UNLESS(Size() > 0 && Position < (ui64)Column(0).length());
+        for (auto&& i : *Columns) {
+            Types.emplace_back(i->type_id());
+        }
     }
 
     template<typename T>
     bool operator == (const TReplaceKeyTemplate<T>& key) const {
         Y_ABORT_UNLESS(Size() == key.Size());
 
-        for (int i = 0; i < Size(); ++i) {
+        for (ui32 i = 0; i < Size(); ++i) {
             auto cmp = CompareColumnValue(i, key, i);
             if (std::is_neq(cmp)) {
                 return false;
@@ -84,7 +90,7 @@ class TReplaceKeyTemplate {
     std::partial_ordering operator <=> (const TReplaceKeyTemplate<T>& key) const {
         Y_ABORT_UNLESS(Size() == key.Size());
 
-        for (int i = 0; i < Size(); ++i) {
+        for (ui32 i = 0; i < Size(); ++i) {
             auto cmp = CompareColumnValue(i, key, i);
             if (std::is_neq(cmp)) {
                 return cmp;
@@ -97,7 +103,7 @@ class TReplaceKeyTemplate {
     std::partial_ordering CompareNotNull(const TReplaceKeyTemplate<T>& key) const {
         Y_ABORT_UNLESS(Size() == key.Size());
 
-        for (int i = 0; i < Size(); ++i) {
+        for (ui32 i = 0; i < Size(); ++i) {
             auto cmp = CompareColumnValueNotNull(i, key, i);
             if (std::is_neq(cmp)) {
                 return cmp;
@@ -107,11 +113,11 @@ class TReplaceKeyTemplate {
     }
 
     template<typename T>
-    std::partial_ordering ComparePartNotNull(const TReplaceKeyTemplate<T>& key, int size) const {
+    std::partial_ordering ComparePartNotNull(const TReplaceKeyTemplate<T>& key, const ui32 size) const {
         Y_ABORT_UNLESS(size <= key.Size());
         Y_ABORT_UNLESS(size <= Size());
 
-        for (int i = 0; i < size; ++i) {
+        for (ui32 i = 0; i < size; ++i) {
             auto cmp = CompareColumnValueNotNull(i, key, i);
             if (std::is_neq(cmp)) {
                 return cmp;
@@ -130,11 +136,13 @@ class TReplaceKeyTemplate {
     template<typename T>
     std::partial_ordering CompareColumnValue(int column, const TReplaceKeyTemplate<T>& key, int keyColumn) const {
         Y_DEBUG_ABORT_UNLESS(Column(column).type_id() == key.Column(keyColumn).type_id());
+        Y_DEBUG_ABORT_UNLESS(Types[column] == Column(column).type_id());
+        Y_DEBUG_ABORT_UNLESS(Types.size() == Size());
 
-        return TComparator::TypedCompare<false>(Column(column), Position, key.Column(keyColumn), key.Position);
+        return TComparator::ConcreteTypedCompare<false>(Types[column], Column(column), Position, key.Column(keyColumn), key.Position);
     }
 
-    int Size() const {
+    ui64 Size() const {
         Y_DEBUG_ABORT_UNLESS(Columns);
         return Columns->size();
     }
@@ -166,7 +174,7 @@ class TReplaceKeyTemplate {
 
     template<typename T = TArrayVecPtr> requires IsOwning
     std::shared_ptr<arrow::RecordBatch> RestoreBatch(const std::shared_ptr<arrow::Schema>& schema) const {
-        AFL_VERIFY(Size() && Size() == schema->num_fields())("columns", DebugString())("schema", JoinSeq(",", schema->field_names()));
+        AFL_VERIFY(Size() && Size() == (ui32)schema->num_fields())("columns", DebugString())("schema", JoinSeq(",", schema->field_names()));
         const auto& columns = *Columns;
         return arrow::RecordBatch::Make(schema, columns[0]->length(), columns);
     }
@@ -220,6 +228,7 @@ class TReplaceKeyTemplate {
 
 private:
     TArrayVecPtr Columns = nullptr;
+    std::vector<arrow::Type::type> Types;
     ui64 Position = 0;
 
 };

ydb/core/formats/arrow/switch/compare.h

@@ -7,8 +7,7 @@ namespace NKikimr::NArrow {
 class TComparator {
 public:
     template <bool notNull>
-    static std::partial_ordering TypedCompare(const arrow::Array& lhs, const int lpos, const arrow::Array& rhs, const int rpos) {
-        arrow::Type::type typeId = lhs.type_id();
+    static std::partial_ordering ConcreteTypedCompare(const arrow::Type::type typeId, const arrow::Array& lhs, const int lpos, const arrow::Array& rhs, const int rpos) {
         switch (typeId) {
             case arrow::Type::NA:
             case arrow::Type::BOOL:
@@ -73,6 +72,11 @@ class TComparator {
         return std::partial_ordering::equivalent;
     }
 
+    template <bool notNull>
+    static std::partial_ordering TypedCompare(const arrow::Array& lhs, const int lpos, const arrow::Array& rhs, const int rpos) {
+        return ConcreteTypedCompare<notNull>(lhs.type_id(), lhs, lpos, rhs, rpos);
+    }
+
     template <typename T, bool notNull>
     static std::partial_ordering CompareView(const arrow::Array& lhs, int lpos, const arrow::Array& rhs, int rpos) {
         auto& left = static_cast<const T&>(lhs);