Skip to content

Commit 465faae

Browse files
ssmikessmike
authored
Prefer indices aligned with order-by-limit (#10589)
1 parent ccc35de commit 465faae

File tree

6 files changed

+97
-50
lines changed

6 files changed

+97
-50
lines changed

‎ydb/core/kqp/opt/kqp_opt.cpp‎

Lines changed: 43 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -110,6 +110,49 @@ TKqpTable BuildTableMeta(const TKikimrTableDescription& tableDesc, const TPositi
110110
return BuildTableMeta(*tableDesc.Metadata, pos, ctx);
111111
}
112112

113+
bool IsSortKeyPrimary(const NYql::NNodes::TCoLambda& keySelector, const NYql::TKikimrTableDescription& tableDesc,
114+
const TMaybe<THashSet<TStringBuf>>& passthroughFields)
115+
{
116+
auto checkKey = [keySelector, &tableDesc, &passthroughFields] (NYql::NNodes::TExprBase key, ui32 index) {
117+
if (!key.Maybe<TCoMember>()) {
118+
return false;
119+
}
120+
121+
auto member = key.Cast<TCoMember>();
122+
if (member.Struct().Raw() != keySelector.Args().Arg(0).Raw()) {
123+
return false;
124+
}
125+
126+
auto column = TString(member.Name().Value());
127+
auto columnIndex = tableDesc.GetKeyColumnIndex(column);
128+
if (!columnIndex || *columnIndex != index) {
129+
return false;
130+
}
131+
132+
if (passthroughFields && !passthroughFields->contains(column)) {
133+
return false;
134+
}
135+
136+
return true;
137+
};
138+
139+
auto lambdaBody = keySelector.Body();
140+
if (auto maybeTuple = lambdaBody.Maybe<TExprList>()) {
141+
auto tuple = maybeTuple.Cast();
142+
for (size_t i = 0; i < tuple.Size(); ++i) {
143+
if (!checkKey(tuple.Item(i), i)) {
144+
return false;
145+
}
146+
}
147+
} else {
148+
if (!checkKey(lambdaBody, 0)) {
149+
return false;
150+
}
151+
}
152+
153+
return true;
154+
}
155+
113156
bool IsBuiltEffect(const TExprBase& effect) {
114157
// Stage with effect output
115158
if (effect.Maybe<TDqOutput>()) {

‎ydb/core/kqp/opt/kqp_opt_impl.h‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -62,4 +62,7 @@ TVector<std::pair<NYql::TExprNode::TPtr, const NYql::TIndexDescription*>> BuildS
6262

6363
bool IsBuiltEffect(const NYql::NNodes::TExprBase& effect);
6464

65+
bool IsSortKeyPrimary(const NYql::NNodes::TCoLambda& keySelector, const NYql::TKikimrTableDescription& tableDesc,
66+
const TMaybe<THashSet<TStringBuf>>& passthroughFields = {});
67+
6568
} // namespace NKikimr::NKqp::NOpt

‎ydb/core/kqp/opt/logical/kqp_opt_log_ranges_predext.cpp‎

Lines changed: 28 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -177,6 +177,22 @@ TMaybeNode<TExprBase> TryBuildTrivialReadTable(TCoFlatMap& flatmap, TKqlReadTabl
177177
.Done();
178178
}
179179

180+
TMaybeNode<TCoLambda> ExtractTopSortKeySelector(TExprBase node, const NYql::TParentsMap& parentsMap) {
181+
auto it = parentsMap.find(node.Raw());
182+
if (it != parentsMap.end()) {
183+
if (it->second.size() != 1) {
184+
return {};
185+
}
186+
for (auto* node : it->second) {
187+
if (TCoTopSort::Match(node)) {
188+
TCoTopSort topSort(node);
189+
return topSort.KeySelectorLambda();
190+
}
191+
}
192+
}
193+
return {};
194+
}
195+
180196
} // namespace
181197

182198
TExprBase KqpPushExtractedPredicateToReadTable(TExprBase node, TExprContext& ctx, const TKqpOptimizeContext& kqpCtx,
@@ -269,7 +285,7 @@ TExprBase KqpPushExtractedPredicateToReadTable(TExprBase node, TExprContext& ctx
269285
YQL_ENSURE(prepareSuccess);
270286

271287
if (!indexName.IsValid() && !readSettings.ForcePrimary && kqpCtx.Config->IndexAutoChooserMode != NKikimrConfig::TTableServiceConfig_EIndexAutoChooseMode_DISABLED) {
272-
using TIndexComparisonKey = std::tuple<bool, size_t, bool, size_t, bool>;
288+
using TIndexComparisonKey = std::tuple<bool, bool, size_t, bool, size_t, bool>;
273289
auto calcNeedsJoin = [&] (const TKikimrTableMetadataPtr& keyTable) -> bool {
274290
bool needsJoin = false;
275291
for (auto&& column : read.Columns()) {
@@ -280,8 +296,16 @@ TExprBase KqpPushExtractedPredicateToReadTable(TExprBase node, TExprContext& ctx
280296
return needsJoin;
281297
};
282298

283-
auto calcKey = [&](NYql::IPredicateRangeExtractor::TBuildResult buildResult, size_t descriptionKeyColumns, bool needsJoin) -> TIndexComparisonKey {
299+
auto keySelector = ExtractTopSortKeySelector(flatmap, parentsMap);
300+
301+
auto calcKey = [&](
302+
NYql::IPredicateRangeExtractor::TBuildResult buildResult,
303+
size_t descriptionKeyColumns,
304+
bool needsJoin,
305+
const NYql::TKikimrTableDescription & tableDesc) -> TIndexComparisonKey
306+
{
284307
return std::make_tuple(
308+
keySelector.IsValid() && IsSortKeyPrimary(keySelector.Cast(), tableDesc),
285309
buildResult.PointPrefixLen >= descriptionKeyColumns,
286310
buildResult.PointPrefixLen >= descriptionKeyColumns ? 0 : buildResult.PointPrefixLen,
287311
buildResult.UsedPrefixLen >= descriptionKeyColumns,
@@ -293,7 +317,7 @@ TExprBase KqpPushExtractedPredicateToReadTable(TExprBase node, TExprContext& ctx
293317
auto primaryBuildResult = extractor->BuildComputeNode(mainTableDesc.Metadata->KeyColumnNames, ctx, typesCtx);
294318

295319
if (primaryBuildResult.PointPrefixLen < mainTableDesc.Metadata->KeyColumnNames.size()) {
296-
auto maxKey = calcKey(primaryBuildResult, mainTableDesc.Metadata->KeyColumnNames.size(), false);
320+
auto maxKey = calcKey(primaryBuildResult, mainTableDesc.Metadata->KeyColumnNames.size(), false, mainTableDesc);
297321
for (auto& index : mainTableDesc.Metadata->Indexes) {
298322
if (index.Type != TIndexDescription::EType::GlobalAsync) {
299323
auto& tableDesc = kqpCtx.Tables->ExistingTable(kqpCtx.Cluster, mainTableDesc.Metadata->GetIndexMetadata(TString(index.Name)).first->Name);
@@ -307,7 +331,7 @@ TExprBase KqpPushExtractedPredicateToReadTable(TExprBase node, TExprContext& ctx
307331
continue;
308332
}
309333

310-
auto key = calcKey(buildResult, index.KeyColumns.size(), needsJoin);
334+
auto key = calcKey(buildResult, index.KeyColumns.size(), needsJoin, tableDesc);
311335
if (key > maxKey) {
312336
maxKey = key;
313337
chosenIndex = index.Name;

‎ydb/core/kqp/opt/physical/kqp_opt_phy_helpers.cpp‎

Lines changed: 0 additions & 43 deletions
Original file line numberDiff line numberDiff line change
@@ -170,49 +170,6 @@ NYql::NNodes::TDqStage ReplaceTableSourceSettings(NYql::NNodes::TDqStage stage,
170170
.Done();
171171
}
172172

173-
bool IsSortKeyPrimary(const NYql::NNodes::TCoLambda& keySelector, const NYql::TKikimrTableDescription& tableDesc,
174-
const TMaybe<THashSet<TStringBuf>>& passthroughFields)
175-
{
176-
auto checkKey = [keySelector, &tableDesc, &passthroughFields] (NYql::NNodes::TExprBase key, ui32 index) {
177-
if (!key.Maybe<TCoMember>()) {
178-
return false;
179-
}
180-
181-
auto member = key.Cast<TCoMember>();
182-
if (member.Struct().Raw() != keySelector.Args().Arg(0).Raw()) {
183-
return false;
184-
}
185-
186-
auto column = TString(member.Name().Value());
187-
auto columnIndex = tableDesc.GetKeyColumnIndex(column);
188-
if (!columnIndex || *columnIndex != index) {
189-
return false;
190-
}
191-
192-
if (passthroughFields && !passthroughFields->contains(column)) {
193-
return false;
194-
}
195-
196-
return true;
197-
};
198-
199-
auto lambdaBody = keySelector.Body();
200-
if (auto maybeTuple = lambdaBody.Maybe<TExprList>()) {
201-
auto tuple = maybeTuple.Cast();
202-
for (size_t i = 0; i < tuple.Size(); ++i) {
203-
if (!checkKey(tuple.Item(i), i)) {
204-
return false;
205-
}
206-
}
207-
} else {
208-
if (!checkKey(lambdaBody, 0)) {
209-
return false;
210-
}
211-
}
212-
213-
return true;
214-
}
215-
216173
ESortDirection GetSortDirection(const NYql::NNodes::TExprBase& sortDirections) {
217174
auto getDirection = [] (TExprBase expr) -> ESortDirection {
218175
if (!expr.Maybe<TCoBool>()) {

‎ydb/core/kqp/opt/physical/kqp_opt_phy_impl.h‎

Lines changed: 0 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -28,9 +28,6 @@ NYql::NNodes::TDqStage ReplaceStageArg(NYql::NNodes::TDqStage stage, size_t inpu
2828
NYql::NNodes::TDqStage ReplaceTableSourceSettings(NYql::NNodes::TDqStage stage, size_t inputIndex,
2929
NYql::NNodes::TKqpReadRangesSourceSettings settings, NYql::TExprContext& ctx);
3030

31-
bool IsSortKeyPrimary(const NYql::NNodes::TCoLambda& keySelector, const NYql::TKikimrTableDescription& tableDesc,
32-
const TMaybe<THashSet<TStringBuf>>& passthroughFields = {});
33-
3431
enum ESortDirection : ui32 {
3532
None = 0,
3633
Forward = 1,

‎ydb/core/kqp/ut/opt/kqp_ne_ut.cpp‎

Lines changed: 23 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -4059,6 +4059,29 @@ Y_UNIT_TEST_SUITE(KqpNewEngine) {
40594059
AssertTableReads(result, "/Root/SecondaryKeys/Index/indexImplTable", 1);
40604060
}
40614061

4062+
Y_UNIT_TEST(AutoChooseIndexOrderByLimit) {
4063+
TKikimrSettings settings;
4064+
NKikimrConfig::TAppConfig appConfig;
4065+
appConfig.MutableTableServiceConfig()->SetIndexAutoChooseMode(NKikimrConfig::TTableServiceConfig_EIndexAutoChooseMode_ONLY_POINTS);
4066+
settings.SetAppConfig(appConfig);
4067+
4068+
TKikimrRunner kikimr(settings);
4069+
4070+
auto db = kikimr.GetTableClient();
4071+
auto session = db.CreateSession().GetValueSync().GetSession();
4072+
CreateSampleTablesWithIndex(session);
4073+
4074+
NYdb::NTable::TExecDataQuerySettings querySettings;
4075+
querySettings.CollectQueryStats(ECollectQueryStatsMode::Profile);
4076+
4077+
auto result = session.ExecuteDataQuery(R"(
4078+
--!syntax_v1
4079+
SELECT Fk, Key FROM `/Root/SecondaryKeys` WHERE Fk = 1 ORDER BY Key DESC LIMIT 1;
4080+
)", TTxControl::BeginTx(TTxSettings::SerializableRW()), querySettings).GetValueSync();
4081+
AssertSuccessResult(result);
4082+
AssertTableReads(result, "/Root/SecondaryKeys/Index/indexImplTable", 0);
4083+
}
4084+
40624085
Y_UNIT_TEST(MultipleBroadcastJoin) {
40634086
TKikimrSettings kisettings;
40644087
NKikimrConfig::TAppConfig appConfig;

0 commit comments

Comments
 (0)