Skip to content

Commit

Permalink
[Internal] Query: Removes unneeded check on supported query features (#…
Browse files Browse the repository at this point in the history
…4868)

This PR

-Upgrades to the 3.37.1 version of the direct package
-Adds additional testing on ORDER BY RANK RRF with VectorDistance in
HybridSearchQueryTests SanityTests
-Removes an unneeded (and problematic) check for unsupported features,
which is not needed by the SDK. Previously, this code assumed there
would always be a queryInfo object, however with hybrid search, this may
not always be the case. But in addition, this check was not needed, so
it is removed.

Please delete options that are not relevant.

- [] Bug fix (non-breaking change which fixes an issue)
  • Loading branch information
sc978345 authored and kundadebdatta committed Nov 10, 2024
1 parent e665fa8 commit 2efd1a9
Show file tree
Hide file tree
Showing 4 changed files with 52 additions and 273 deletions.
262 changes: 0 additions & 262 deletions Microsoft.Azure.Cosmos/src/Query/Core/QueryPlan/QueryPlanHandler.cs
Original file line number Diff line number Diff line change
Expand Up @@ -27,7 +27,6 @@ public async Task<TryCatch<PartitionedQueryExecutionInfo>> TryGetQueryPlanAsync(
Documents.ResourceType resourceType,
PartitionKeyDefinition partitionKeyDefinition,
VectorEmbeddingPolicy vectorEmbeddingPolicy,
QueryFeatures supportedQueryFeatures,
bool hasLogicalPartitionKey,
bool useSystemPrefix,
GeospatialType geospatialType,
Expand Down Expand Up @@ -59,63 +58,9 @@ public async Task<TryCatch<PartitionedQueryExecutionInfo>> TryGetQueryPlanAsync(
return tryGetQueryInfo;
}

if (QueryPlanExceptionFactory.TryGetUnsupportedException(
tryGetQueryInfo.Result.QueryInfo,
supportedQueryFeatures,
out Exception queryPlanHandlerException))
{
return TryCatch<PartitionedQueryExecutionInfo>.FromException(queryPlanHandlerException);
}

return tryGetQueryInfo;
}

/// <summary>
/// Used in the compute gateway to support legacy gateways query execution pattern.
/// </summary>
public async Task<TryCatch<(PartitionedQueryExecutionInfo queryPlan, bool supported)>> TryGetQueryInfoAndIfSupportedAsync(
QueryFeatures supportedQueryFeatures,
SqlQuerySpec sqlQuerySpec,
Documents.ResourceType resourceType,
PartitionKeyDefinition partitionKeyDefinition,
VectorEmbeddingPolicy vectorEmbeddingPolicy,
bool hasLogicalPartitionKey,
bool useSystemPrefix,
GeospatialType geospatialType,
CancellationToken cancellationToken = default)
{
if (sqlQuerySpec == null)
{
throw new ArgumentNullException(nameof(sqlQuerySpec));
}

if (partitionKeyDefinition == null)
{
throw new ArgumentNullException(nameof(partitionKeyDefinition));
}

cancellationToken.ThrowIfCancellationRequested();

TryCatch<PartitionedQueryExecutionInfo> tryGetQueryInfo = await this.TryGetQueryInfoAsync(
sqlQuerySpec,
resourceType,
partitionKeyDefinition,
vectorEmbeddingPolicy,
hasLogicalPartitionKey,
useSystemPrefix,
geospatialType,
cancellationToken);
if (tryGetQueryInfo.Failed)
{
return TryCatch<(PartitionedQueryExecutionInfo, bool)>.FromException(tryGetQueryInfo.Exception);
}

QueryFeatures neededQueryFeatures = QueryPlanSupportChecker.GetNeededQueryFeatures(
tryGetQueryInfo.Result.QueryInfo,
supportedQueryFeatures);
return TryCatch<(PartitionedQueryExecutionInfo, bool)>.FromResult((tryGetQueryInfo.Result, neededQueryFeatures == QueryFeatures.None));
}

private Task<TryCatch<PartitionedQueryExecutionInfo>> TryGetQueryInfoAsync(
SqlQuerySpec sqlQuerySpec,
Documents.ResourceType resourceType,
Expand All @@ -142,212 +87,5 @@ private Task<TryCatch<PartitionedQueryExecutionInfo>> TryGetQueryInfoAsync(
geospatialType: geospatialType,
cancellationToken: cancellationToken);
}

private static class QueryPlanSupportChecker
{
public static QueryFeatures GetNeededQueryFeatures(
QueryInfo queryInfo,
QueryFeatures supportedQueryFeatures)
{
QueryFeatures neededQueryFeatures = QueryFeatures.None;
neededQueryFeatures |= QueryPlanSupportChecker.GetNeededQueryFeaturesIfAggregateQuery(queryInfo, supportedQueryFeatures);
neededQueryFeatures |= QueryPlanSupportChecker.GetNeededQueryFeaturesIfDistinctQuery(queryInfo, supportedQueryFeatures);
neededQueryFeatures |= QueryPlanSupportChecker.GetNeedQueryFeaturesIfGroupByQuery(queryInfo, supportedQueryFeatures);
neededQueryFeatures |= QueryPlanSupportChecker.GetNeededQueryFeaturesIfOffsetLimitQuery(queryInfo, supportedQueryFeatures);
neededQueryFeatures |= QueryPlanSupportChecker.GetNeededQueryFeaturesIfOrderByQuery(queryInfo, supportedQueryFeatures);
neededQueryFeatures |= QueryPlanSupportChecker.GetNeededQueryFeaturesIfTopQuery(queryInfo, supportedQueryFeatures);

return neededQueryFeatures;
}

private static QueryFeatures GetNeededQueryFeaturesIfAggregateQuery(
QueryInfo queryInfo,
QueryFeatures supportedQueryFeatures)
{
QueryFeatures neededQueryFeatures = QueryFeatures.None;
if (queryInfo.HasAggregates)
{
bool isSingleAggregate = (queryInfo.Aggregates.Count == 1)
|| (queryInfo.GroupByAliasToAggregateType.Values.Where(aggregateOperator => aggregateOperator.HasValue).Count() == 1);
if (isSingleAggregate)
{
if (queryInfo.HasSelectValue)
{
if (!supportedQueryFeatures.HasFlag(QueryFeatures.Aggregate))
{
neededQueryFeatures |= QueryFeatures.Aggregate;
}
}
else
{
if (!supportedQueryFeatures.HasFlag(QueryFeatures.NonValueAggregate))
{
neededQueryFeatures |= QueryFeatures.NonValueAggregate;
}
}
}
else
{
if (!supportedQueryFeatures.HasFlag(QueryFeatures.NonValueAggregate))
{
neededQueryFeatures |= QueryFeatures.NonValueAggregate;
}

if (!supportedQueryFeatures.HasFlag(QueryFeatures.MultipleAggregates))
{
neededQueryFeatures |= QueryFeatures.MultipleAggregates;
}
}
}

return neededQueryFeatures;
}

private static QueryFeatures GetNeededQueryFeaturesIfDistinctQuery(
QueryInfo queryInfo,
QueryFeatures supportedQueryFeatures)
{
QueryFeatures neededQueryFeatures = QueryFeatures.None;
if (queryInfo.HasDistinct)
{
if (!supportedQueryFeatures.HasFlag(QueryFeatures.Distinct))
{
neededQueryFeatures |= QueryFeatures.Distinct;
}
}

return neededQueryFeatures;
}

private static QueryFeatures GetNeededQueryFeaturesIfTopQuery(
QueryInfo queryInfo,
QueryFeatures supportedQueryFeatures)
{
QueryFeatures neededQueryFeatures = QueryFeatures.None;
if (queryInfo.HasTop)
{
if (!supportedQueryFeatures.HasFlag(QueryFeatures.Top))
{
neededQueryFeatures |= QueryFeatures.Top;
}
}

return neededQueryFeatures;
}

private static QueryFeatures GetNeededQueryFeaturesIfOffsetLimitQuery(
QueryInfo queryInfo,
QueryFeatures supportedQueryFeatures)
{
QueryFeatures neededQueryFeatures = QueryFeatures.None;
if (queryInfo.HasLimit || queryInfo.HasOffset)
{
if (!supportedQueryFeatures.HasFlag(QueryFeatures.OffsetAndLimit))
{
neededQueryFeatures |= QueryFeatures.OffsetAndLimit;
}
}

return neededQueryFeatures;
}

private static QueryFeatures GetNeedQueryFeaturesIfGroupByQuery(
QueryInfo queryInfo,
QueryFeatures supportedQueryFeatures)
{
QueryFeatures neededQueryFeatures = QueryFeatures.None;
if (queryInfo.HasGroupBy)
{
if (!supportedQueryFeatures.HasFlag(QueryFeatures.GroupBy))
{
neededQueryFeatures |= QueryFeatures.GroupBy;
}
}

return neededQueryFeatures;
}

private static QueryFeatures GetNeededQueryFeaturesIfOrderByQuery(
QueryInfo queryInfo,
QueryFeatures supportedQueryFeatures)
{
QueryFeatures neededQueryFeatures = QueryFeatures.None;
if (queryInfo.HasOrderBy)
{
if (queryInfo.OrderByExpressions.Count == 1)
{
if (!supportedQueryFeatures.HasFlag(QueryFeatures.OrderBy))
{
neededQueryFeatures |= QueryFeatures.OrderBy;
}
}
else
{
if (!supportedQueryFeatures.HasFlag(QueryFeatures.MultipleOrderBy))
{
neededQueryFeatures |= QueryFeatures.MultipleOrderBy;
}
}
}

return neededQueryFeatures;
}
}

private static class QueryPlanExceptionFactory
{
private static readonly IReadOnlyList<QueryFeatures> QueryFeatureList = (QueryFeatures[])Enum.GetValues(typeof(QueryFeatures));
private static readonly ReadOnlyDictionary<QueryFeatures, ArgumentException> FeatureToUnsupportedException = new ReadOnlyDictionary<QueryFeatures, ArgumentException>(
QueryFeatureList
.ToDictionary(
x => x,
x => new ArgumentException(QueryPlanExceptionFactory.FormatExceptionMessage(x.ToString()))));

public static bool TryGetUnsupportedException(
QueryInfo queryInfo,
QueryFeatures supportedQueryFeatures,
out Exception queryPlanHandlerException)
{
QueryFeatures neededQueryFeatures = QueryPlanSupportChecker.GetNeededQueryFeatures(
queryInfo,
supportedQueryFeatures);
if (neededQueryFeatures != QueryFeatures.None)
{
List<Exception> queryPlanHandlerExceptions = new List<Exception>();
foreach (QueryFeatures queryFeature in QueryPlanExceptionFactory.QueryFeatureList)
{
if ((neededQueryFeatures & queryFeature) == queryFeature)
{
Exception unsupportedFeatureException = QueryPlanExceptionFactory.FeatureToUnsupportedException[queryFeature];
queryPlanHandlerExceptions.Add(unsupportedFeatureException);
}
}

queryPlanHandlerException = new QueryPlanHandlerException(queryPlanHandlerExceptions);
return true;
}

queryPlanHandlerException = default;
return false;
}

private static string FormatExceptionMessage(string feature)
{
return $"Query contained {feature}, which the calling client does not support.";
}

private sealed class QueryPlanHandlerException : AggregateException
{
private const string QueryContainsUnsupportedFeaturesExceptionMessage = "Query contains 1 or more unsupported features. Upgrade your SDK to a version that does support the requested features:";
public QueryPlanHandlerException(IEnumerable<Exception> innerExceptions)
: base(
QueryContainsUnsupportedFeaturesExceptionMessage
+ Environment.NewLine
+ string.Join(Environment.NewLine, innerExceptions.Select(innerException => innerException.Message)),
innerExceptions)
{
}
}
}
}
}
Original file line number Diff line number Diff line change
Expand Up @@ -86,7 +86,6 @@ public static async Task<PartitionedQueryExecutionInfo> GetQueryPlanWithServiceI
resourceType,
partitionKeyDefinition,
vectorEmbeddingPolicy,
QueryPlanRetriever.SupportedQueryFeatures,
hasLogicalPartitionKey,
useSystemPrefix,
geospatialType,
Expand Down
Loading

0 comments on commit 2efd1a9

Please sign in to comment.