Merge b740c1b into d7e5925

ydb-platform · Jan 5, 2025 · f9ba26d · f9ba26d
2 parents d7e5925 + b740c1b
commit f9ba26d
Show file tree

Hide file tree

Showing 2 changed files with 33 additions and 12 deletions.
diff --git a/ydb/core/kqp/runtime/kqp_stream_lookup_worker.cpp b/ydb/core/kqp/runtime/kqp_stream_lookup_worker.cpp
@@ -19,10 +19,31 @@ std::vector<std::pair<ui64, TOwnedTableRange>> GetRangePartitioning(const TKqpSt
 
     YQL_ENSURE(partitionInfo);
 
+    size_t idxStart = 0;
+    size_t idxFinish = partitionInfo->size();
+    while ((idxFinish - idxStart) > 1) {
+        size_t idxCur = (idxFinish + idxStart) / 2;
+        int cmp = CompareTypedCellVectors
+            ( (*partitionInfo)[idxCur].Range->EndKeyPrefix.GetCells().data(),
+              range.From.data(),
+              keyColumnTypes.data(),
+              keyColumnTypes.size() );
+        if (cmp < 0) {
+            idxStart = idxCur + 1;
+        } else {
+            idxFinish = idxCur;
+        }
+    }
+    if (idxStart > 0) {
+        idxStart -= 1;
+    } else if (idxStart >= partitionInfo->size()) {
+        idxStart = partitionInfo->size() - 1;
+    }
+
     std::vector<TCell> minusInf(keyColumnTypes.size());
 
     std::vector<std::pair<ui64, TOwnedTableRange>> rangePartition;
-    for (size_t idx = 0; idx < partitionInfo->size(); ++idx) {
+    for (size_t idx = idxStart; idx < partitionInfo->size(); ++idx) {
         TTableRange partitionRange{
             idx == 0 ? minusInf : (*partitionInfo)[idx - 1].Range->EndKeyPrefix.GetCells(),
             idx == 0 ? true : !(*partitionInfo)[idx - 1].Range->IsInclusive,
@@ -110,6 +131,12 @@ TKqpStreamLookupWorker::TKqpStreamLookupWorker(NKikimrKqp::TKqpStreamLookupSetti
             column.GetTypeInfo().GetPgTypeMod()
         });
     }
+
+    KeyColumnTypes.resize(KeyColumns.size());
+    for (const auto& [_, columnInfo] : KeyColumns) {
+        YQL_ENSURE(columnInfo.KeyOrder < static_cast<i64>(KeyColumnTypes.size()));
+        KeyColumnTypes[columnInfo.KeyOrder] = columnInfo.PType;
+    }
 }
 
 TKqpStreamLookupWorker::~TKqpStreamLookupWorker() {
@@ -123,16 +150,6 @@ TTableId TKqpStreamLookupWorker::GetTableId() const {
     return TableId;
 }
 
-std::vector<NScheme::TTypeInfo> TKqpStreamLookupWorker::GetKeyColumnTypes() const {
-    std::vector<NScheme::TTypeInfo> keyColumnTypes(KeyColumns.size());
-    for (const auto& [_, columnInfo] : KeyColumns) {
-        YQL_ENSURE(columnInfo.KeyOrder < static_cast<i64>(keyColumnTypes.size()));
-        keyColumnTypes[columnInfo.KeyOrder] = columnInfo.PType;
-    }
-
-    return keyColumnTypes;
-}
-
 class TKqpLookupRows : public TKqpStreamLookupWorker {
 public:
     TKqpLookupRows(NKikimrKqp::TKqpStreamLookupSettings&& settings, const NMiniKQL::TTypeEnvironment& typeEnv,

diff --git a/ydb/core/kqp/runtime/kqp_stream_lookup_worker.h b/ydb/core/kqp/runtime/kqp_stream_lookup_worker.h
@@ -51,7 +51,10 @@ class TKqpStreamLookupWorker {
 
     virtual std::string GetTablePath() const;
     virtual TTableId GetTableId() const;
-    virtual std::vector<NScheme::TTypeInfo> GetKeyColumnTypes() const;
+
+    const std::vector<NScheme::TTypeInfo>& GetKeyColumnTypes() const {
+        return KeyColumnTypes;
+    }
 
     virtual void AddInputRow(NUdf::TUnboxedValue inputRow) = 0;
     virtual std::vector<THolder<TEvDataShard::TEvRead>> RebuildRequest(const ui64& prevReadId, ui32 firstUnprocessedQuery, 
@@ -72,6 +75,7 @@ class TKqpStreamLookupWorker {
     std::unordered_map<TString, TSysTables::TTableColumnInfo> KeyColumns;
     std::vector<TSysTables::TTableColumnInfo*> LookupKeyColumns;
     std::vector<TSysTables::TTableColumnInfo> Columns;
+    std::vector<NScheme::TTypeInfo> KeyColumnTypes;
 };
 
 std::unique_ptr<TKqpStreamLookupWorker> CreateStreamLookupWorker(NKikimrKqp::TKqpStreamLookupSettings&& settings,