Filter stat for ReferenceSearchParam and UriSearchParam table

src/Microsoft.Health.Fhir.SqlServer.UnitTests/Features/Search/SqlServerSearchServiceTests.cs

@@ -298,5 +298,57 @@ public void Model_Property_ReturnsInjectedModel()
             // Assert
             Assert.Same(_model, model);
         }
+
+        public static IEnumerable<object[]> SingleColumnTableData()
+        {
+            yield return new object[] { VLatest.TokenSearchParam.TableName, VLatest.TokenSearchParam.Code.Metadata.Name };
+            yield return new object[] { VLatest.StringSearchParam.TableName, VLatest.StringSearchParam.Text.Metadata.Name };
+            yield return new object[] { VLatest.UriSearchParam.TableName, VLatest.UriSearchParam.Uri.Metadata.Name };
+        }
+
+        [Theory]
+        [MemberData(nameof(SingleColumnTableData))]
+        public void GetKeyColumns_ForSingleColumnTable_ReturnsOnlyExpectedColumn(string tableName, string expectedColumn)
+        {
+            // Act
+            var columns = SqlServerSearchService.GetKeyColumns(tableName);
+
+            // Assert
+            Assert.Contains(expectedColumn, columns);
+            Assert.Single(columns);
+        }
+
+        public static IEnumerable<object[]> TwoColumnTableData()
+        {
+            yield return new object[] { VLatest.DateTimeSearchParam.TableName, VLatest.DateTimeSearchParam.StartDateTime.Metadata.Name, VLatest.DateTimeSearchParam.EndDateTime.Metadata.Name };
+            yield return new object[] { VLatest.NumberSearchParam.TableName, VLatest.NumberSearchParam.LowValue.Metadata.Name, VLatest.NumberSearchParam.HighValue.Metadata.Name };
+            yield return new object[] { VLatest.QuantitySearchParam.TableName, VLatest.QuantitySearchParam.LowValue.Metadata.Name, VLatest.QuantitySearchParam.HighValue.Metadata.Name };
+            yield return new object[] { VLatest.ReferenceSearchParam.TableName, VLatest.ReferenceSearchParam.ReferenceResourceId.Metadata.Name, VLatest.ReferenceSearchParam.ReferenceResourceTypeId.Metadata.Name };
+        }
+
+        [Theory]
+        [MemberData(nameof(TwoColumnTableData))]
+        public void GetKeyColumns_ForTwoColumnTable_ReturnsBothExpectedColumns(string tableName, string column1, string column2)
+        {
+            // Act
+            var columns = SqlServerSearchService.GetKeyColumns(tableName);
+
+            // Assert
+            Assert.Contains(column1, columns);
+            Assert.Contains(column2, columns);
+            Assert.Equal(2, columns.Count);
+        }
+
+        [Theory]
+        [InlineData("dbo.UnknownTable")] // NotExists expressions resolve to no known table — stats must be skipped
+        [InlineData(null)]
+        public void GetKeyColumns_ForUnknownOrNullTable_ReturnsEmptySet(string tableName)
+        {
+            // Act
+            var columns = SqlServerSearchService.GetKeyColumns(tableName);
+
+            // Assert
+            Assert.Empty(columns);
+        }
     }
 }

src/Microsoft.Health.Fhir.SqlServer/Features/Search/SqlServerSearchService.cs

@@ -1140,6 +1140,51 @@ internal static string StripAllWhitespace(string text)
         internal static string StripDboSchemaPrefix(string procName) =>
             procName?.Replace("dbo.", string.Empty, StringComparison.OrdinalIgnoreCase);
 
+        /// <summary>
+        /// Returns the column names used for filtered statistics on the given search parameter table.
+        /// Only tables that benefit from per-resource-type filtered statistics are included.
+        /// </summary>
+        /// <param name="table">The fully-qualified table name (e.g. <c>dbo.TokenSearchParam</c>).</param>
+        /// <returns>The set of column names, or an empty set when the table does not support filtered stats.</returns>
+        internal static HashSet<string> GetKeyColumns(string table)
+        {
+            var results = new HashSet<string>();
+            if (table == VLatest.StringSearchParam.TableName)
+            {
+                results.Add(VLatest.StringSearchParam.Text.Metadata.Name);
+            }
+            else if (table == VLatest.TokenSearchParam.TableName)
+            {
+                results.Add(VLatest.TokenSearchParam.Code.Metadata.Name);
+            }
+            else if (table == VLatest.DateTimeSearchParam.TableName)
+            {
+                results.Add(VLatest.DateTimeSearchParam.StartDateTime.Metadata.Name);
+                results.Add(VLatest.DateTimeSearchParam.EndDateTime.Metadata.Name);
+            }
+            else if (table == VLatest.NumberSearchParam.TableName)
+            {
+                results.Add(VLatest.NumberSearchParam.LowValue.Metadata.Name);
+                results.Add(VLatest.NumberSearchParam.HighValue.Metadata.Name);
+            }
+            else if (table == VLatest.QuantitySearchParam.TableName)
+            {
+                results.Add(VLatest.QuantitySearchParam.LowValue.Metadata.Name);
+                results.Add(VLatest.QuantitySearchParam.HighValue.Metadata.Name);
+            }
+            else if (table == VLatest.ReferenceSearchParam.TableName)
+            {
+                results.Add(VLatest.ReferenceSearchParam.ReferenceResourceId.Metadata.Name);
+                results.Add(VLatest.ReferenceSearchParam.ReferenceResourceTypeId.Metadata.Name);
+            }
+            else if (table == VLatest.UriSearchParam.TableName)
+            {
+                results.Add(VLatest.UriSearchParam.Uri.Metadata.Name);
+            }
+
+            return results;
+        }
+
         private async Task LogQueryStoreByTextAsync(
             string queryText,
             bool isStoredProcedure,
@@ -2542,35 +2587,7 @@ private static void CollectResourceTypesFromExpression(Expression expression, Sq
                 }
             }
 
-            private static HashSet<string> GetKeyColumns(string table)
-            {
-                var results = new HashSet<string>();
-                if (table == VLatest.StringSearchParam.TableName)
-                {
-                    results.Add(VLatest.StringSearchParam.Text.Metadata.Name);
-                }
-                else if (table == VLatest.TokenSearchParam.TableName)
-                {
-                    results.Add(VLatest.TokenSearchParam.Code.Metadata.Name);
-                }
-                else if (table == VLatest.DateTimeSearchParam.TableName)
-                {
-                    results.Add(VLatest.DateTimeSearchParam.StartDateTime.Metadata.Name);
-                    results.Add(VLatest.DateTimeSearchParam.EndDateTime.Metadata.Name);
-                }
-                else if (table == VLatest.NumberSearchParam.TableName)
-                {
-                    results.Add(VLatest.NumberSearchParam.LowValue.Metadata.Name);
-                    results.Add(VLatest.NumberSearchParam.HighValue.Metadata.Name);
-                }
-                else if (table == VLatest.QuantitySearchParam.TableName)
-                {
-                    results.Add(VLatest.QuantitySearchParam.LowValue.Metadata.Name);
-                    results.Add(VLatest.QuantitySearchParam.HighValue.Metadata.Name);
-                }
-
-                return results;
-            }
+            private static HashSet<string> GetKeyColumns(string table) => SqlServerSearchService.GetKeyColumns(table);
 
             private async Task Create(string tableName, string columnName, short resourceTypeId, short searchParamId, ISqlRetryService sqlRetryService, ILogger<SqlServerSearchService> logger, CancellationToken cancel)
             {

test/Microsoft.Health.Fhir.Shared.Tests.Integration/Persistence/SqlServerCreateStatsTests.cs

@@ -100,5 +100,130 @@ public async Task GivenSearchForResearchStudyByFocusAndDateWithResearchSubject_S
                       && _.SearchParamId == model.GetSearchParamId(new Uri("http://hl7.org/fhir/SearchParameter/ResearchSubject-status")));
             }
         }
+
+        [Fact]
+        public async Task GivenSearchByReferenceParam_NormalPath_StatsAreCreated()
+        {
+            // Arrange — DiagnosticReport.subject is a reference search parameter stored in dbo.ReferenceSearchParam
+            const string resourceType = "DiagnosticReport";
+            const string searchParamUrl = "http://hl7.org/fhir/SearchParameter/DiagnosticReport-subject";
+            var query = new[] { Tuple.Create("subject", "Patient/test-patient-1") };
+            var sqlSearchService = (SqlServerSearchService)_fixture.SearchService;
+            short resourceTypeId = sqlSearchService.Model.GetResourceTypeId(resourceType);
+            short searchParamId = sqlSearchService.Model.GetSearchParamId(new Uri(searchParamUrl));
+
+            // Capture stats before the search so we can assert on the delta.
+            var cacheBefore = SqlServerSearchService.GetStatsFromCache().ToList();
+            var databaseBefore = (await sqlSearchService.GetStatsFromDatabase(CancellationToken.None)).ToList();
+
+            // Act
+            await _fixture.SearchService.SearchAsync(resourceType, query, CancellationToken.None);
+
+            // Assert — filtered statistics must be created for the ReferenceResourceId and ReferenceResourceTypeId columns
+            var cacheAfter = SqlServerSearchService.GetStatsFromCache().ToList();
+            var databaseAfter = (await sqlSearchService.GetStatsFromDatabase(CancellationToken.None)).ToList();
+
+            bool MatchesStat((string TableName, string ColumnName, short ResourceTypeId, short SearchParamId) s, string columnName) =>
+                s.TableName == VLatest.ReferenceSearchParam.TableName
+                && s.ColumnName == columnName
+                && s.ResourceTypeId == resourceTypeId
+                && s.SearchParamId == searchParamId;
+
+            // Cache
+            Assert.True(
+                cacheAfter.Count(s => MatchesStat(s, VLatest.ReferenceSearchParam.ReferenceResourceId.Metadata.Name))
+                > cacheBefore.Count(s => MatchesStat(s, VLatest.ReferenceSearchParam.ReferenceResourceId.Metadata.Name)),
+                "Expected the search to add a filtered statistics entry in the cache for ReferenceResourceId.");
+
+            Assert.True(
+                cacheAfter.Count(s => MatchesStat(s, VLatest.ReferenceSearchParam.ReferenceResourceTypeId.Metadata.Name))
+                > cacheBefore.Count(s => MatchesStat(s, VLatest.ReferenceSearchParam.ReferenceResourceTypeId.Metadata.Name)),
+                "Expected the search to add a filtered statistics entry in the cache for ReferenceResourceTypeId.");
+
+            // Database
+            Assert.True(
+                databaseAfter.Count(s => MatchesStat(s, VLatest.ReferenceSearchParam.ReferenceResourceId.Metadata.Name))
+                > databaseBefore.Count(s => MatchesStat(s, VLatest.ReferenceSearchParam.ReferenceResourceId.Metadata.Name)),
+                "Expected the search to add a filtered statistics entry in the database for ReferenceResourceId.");
+
+            Assert.True(
+                databaseAfter.Count(s => MatchesStat(s, VLatest.ReferenceSearchParam.ReferenceResourceTypeId.Metadata.Name))
+                > databaseBefore.Count(s => MatchesStat(s, VLatest.ReferenceSearchParam.ReferenceResourceTypeId.Metadata.Name)),
+                "Expected the search to add a filtered statistics entry in the database for ReferenceResourceTypeId.");
+        }
+
+        [Fact]
+        public async Task GivenSearchByUriParam_NormalPath_StatsAreCreated()
+        {
+            // Arrange — ValueSet.url is a URI search parameter stored in dbo.UriSearchParam
+            const string resourceType = "ValueSet";
+#if Stu3
+            const string searchParamUrl = "http://hl7.org/fhir/SearchParameter/ValueSet-url";
+#elif R5
+            const string searchParamUrl = "http://hl7.org/fhir/SearchParameter/CanonicalResource-url";
+#else
+            const string searchParamUrl = "http://hl7.org/fhir/SearchParameter/conformance-url";
+#endif
+            var query = new[] { Tuple.Create("url", "http://example.com/fhir/ValueSet/test") };
+            var sqlSearchService = (SqlServerSearchService)_fixture.SearchService;
+            short resourceTypeId = sqlSearchService.Model.GetResourceTypeId(resourceType);
+            short searchParamId = sqlSearchService.Model.GetSearchParamId(new Uri(searchParamUrl));
+
+            bool MatchesStat((string TableName, string ColumnName, short ResourceTypeId, short SearchParamId) s) =>
+                s.TableName == VLatest.UriSearchParam.TableName
+                && s.ColumnName == VLatest.UriSearchParam.Uri.Metadata.Name
+                && s.ResourceTypeId == resourceTypeId
+                && s.SearchParamId == searchParamId;
+
+            // Capture stats before the search so we can assert on the delta.
+            var cacheBefore = SqlServerSearchService.GetStatsFromCache().ToList();
+            var databaseBefore = (await sqlSearchService.GetStatsFromDatabase(CancellationToken.None)).ToList();
+
+            // Act
+            await _fixture.SearchService.SearchAsync(resourceType, query, CancellationToken.None);
+
+            // Assert — filtered statistics must be created for the Uri column
+            var cacheAfter = SqlServerSearchService.GetStatsFromCache().ToList();
+            var databaseAfter = (await sqlSearchService.GetStatsFromDatabase(CancellationToken.None)).ToList();
+
+            // Cache
+            Assert.True(
+                cacheAfter.Count(MatchesStat) > cacheBefore.Count(MatchesStat),
+                "Expected the search to add a filtered statistics entry in the cache for Uri.");
+
+            // Database
+            Assert.True(
+                databaseAfter.Count(MatchesStat) > databaseBefore.Count(MatchesStat),
+                "Expected the search to add a filtered statistics entry in the database for Uri.");
+        }
+
+        [Fact]
+        public async Task GivenSearchByReferenceParam_NotExistsPath_NoReferenceStatsAreCreated()
+        {
+            // Arrange — searching with :missing=true produces a NotExists table expression;
+            // the stats pipeline must NOT emit entries for that expression kind.
+            const string resourceType = "DiagnosticReport";
+            var query = new[] { Tuple.Create("subject:missing", "true") };
+
+            // Record the cache before the search so we can detect any new entries added only by this query.
+            var statsBefore = SqlServerSearchService.GetStatsFromCache().ToList();
+
+            // Act
+            await _fixture.SearchService.SearchAsync(resourceType, query, CancellationToken.None);
+
+            var statsAfter = SqlServerSearchService.GetStatsFromCache().ToList();
+
+            // Assert — a NotExists expression must never produce ReferenceSearchParam stats.
+            var newReferenceStats = statsAfter
+                .Where(s => s.TableName == VLatest.ReferenceSearchParam.TableName
+                    && !statsBefore.Any(b =>
+                        b.TableName == s.TableName
+                        && b.ColumnName == s.ColumnName
+                        && b.ResourceTypeId == s.ResourceTypeId
+                        && b.SearchParamId == s.SearchParamId))
+                .ToList();
+
+            Assert.Empty(newReferenceStats);
+        }
     }
 }