Render scalar date day-split unions as a single inline UNION ALL CTE

src/Microsoft.Health.Fhir.Core.UnitTests/Features/Search/Expressions/UnionExpressionTests.cs

@@ -49,5 +49,20 @@ public void GivenUnionExpression_WhenInitializedImproperly_ThrownAnInvalidOperat
                 UnionExpression unionExpressionFail2 = new UnionExpression(UnionOperator.All, new Expression[] { multiaryExpression1 });
             });
         }
+
+        [Fact]
+        public void GivenUnionExpression_WhenCreated_ThenNoSplitMarkerDefaultsToFalseAndCanBeEnabled()
+        {
+            StringExpression expression1 = new StringExpression(StringOperator.Equals, FieldName.String, componentIndex: 0, value: "rush", ignoreCase: true);
+            StringExpression expression2 = new StringExpression(StringOperator.Equals, FieldName.String, componentIndex: 0, value: "2112", ignoreCase: true);
+
+            UnionExpression unionExpression = new UnionExpression(UnionOperator.All, new Expression[] { expression1, expression2 });
+
+            Assert.False(unionExpression.DoNotSplitIntoSeparateCtes);
+
+            unionExpression.DoNotSplitIntoSeparateCtes = true;
+
+            Assert.True(unionExpression.DoNotSplitIntoSeparateCtes);
+        }
     }
 }

src/Microsoft.Health.Fhir.Core/Features/Search/Expressions/ExpressionRewriter.cs

@@ -80,7 +80,16 @@ public virtual Expression VisitMultiary(MultiaryExpression expression, TContext
         public virtual Expression VisitUnion(UnionExpression expression, TContext context)
         {
             IReadOnlyList<Expression> rewrittenExpressions = VisitArray(expression.Expressions, context);
-            return ReferenceEquals(rewrittenExpressions, expression.Expressions) ? expression : new UnionExpression(expression.Operator, rewrittenExpressions);
+
+            if (ReferenceEquals(rewrittenExpressions, expression.Expressions))
+            {
+                return expression;
+            }
+
+            return new UnionExpression(expression.Operator, rewrittenExpressions)
+            {
+                DoNotSplitIntoSeparateCtes = expression.DoNotSplitIntoSeparateCtes,
+            };
         }
 
         public virtual Expression VisitString(StringExpression expression, TContext context)

src/Microsoft.Health.Fhir.Core/Features/Search/Expressions/UnionExpression.cs

@@ -31,6 +31,11 @@ public UnionExpression(UnionOperator unionOperator, IReadOnlyList<Expression> ex
 
         public IReadOnlyList<Expression> Expressions { get; }
 
+        /// <summary>
+        /// Gets or sets a value indicating whether this union must remain in-place and not be split into separate SQL CTE branches.
+        /// </summary>
+        public bool DoNotSplitIntoSeparateCtes { get; set; }
+
         public override TOutput AcceptVisitor<TContext, TOutput>(IExpressionVisitor<TContext, TOutput> visitor, TContext context)
         {
             EnsureArg.IsNotNull(visitor, nameof(visitor));

src/Microsoft.Health.Fhir.SqlServer.UnitTests/Features/Search/Expressions/ScalarTemporalEqualityRewriterTests.cs

@@ -164,6 +164,7 @@ private static void AssertDaySplitUnion(Expression result, DateTimeOffset expect
         {
             var union = Assert.IsType<UnionExpression>(result);
             Assert.Equal(UnionOperator.All, union.Operator);
+            Assert.True(union.DoNotSplitIntoSeparateCtes);
             Assert.Collection(
                 union.Expressions,
                 shortBranch => AssertSearchParameterAnd(

src/Microsoft.Health.Fhir.SqlServer.UnitTests/Features/Search/Expressions/SearchParamTableExpressionExtensionsTests.cs

@@ -0,0 +1,206 @@
+// -------------------------------------------------------------------------------------------------
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License (MIT). See LICENSE in the repo root for license information.
+// -------------------------------------------------------------------------------------------------
+
+using System;
+using System.Collections.Generic;
+using Microsoft.Health.Fhir.Core.Features.Search.Expressions;
+using Microsoft.Health.Fhir.Core.Models;
+using Microsoft.Health.Fhir.SqlServer.Features.Search.Expressions;
+using Microsoft.Health.Fhir.SqlServer.Features.Search.Expressions.Visitors.QueryGenerators;
+using Microsoft.Health.Fhir.Tests.Common;
+using Microsoft.Health.Fhir.ValueSets;
+using Microsoft.Health.Test.Utilities;
+using Xunit;
+
+namespace Microsoft.Health.Fhir.SqlServer.UnitTests.Features.Search.Expressions
+{
+    [Trait(Traits.OwningTeam, OwningTeam.Fhir)]
+    [Trait(Traits.Category, Categories.Search)]
+    public class SearchParamTableExpressionExtensionsTests
+    {
+        private static readonly SearchParameterInfo BirthdateParam = new SearchParameterInfo(
+            "birthdate",
+            "birthdate",
+            SearchParamType.Date,
+            new Uri("http://hl7.org/fhir/SearchParameter/individual-birthdate"),
+            expression: "Patient.birthDate",
+            baseResourceTypes: new[] { "Patient" });
+
+        private static readonly SearchParameterInfo NameParam = new SearchParameterInfo(
+            "name",
+            "name",
+            SearchParamType.String,
+            new Uri("http://hl7.org/fhir/SearchParameter/individual-name"),
+            expression: "Patient.name",
+            baseResourceTypes: new[] { "Patient" });
+
+        public static TheoryData<Expression, bool> UnionDetectionCases => new()
+        {
+            { BuildBareUnion(), true },
+            { Expression.And(BuildBareUnion()), true },
+            { Expression.And(BuildSmartV2Union(), new SearchParameterExpression(NameParam, Expression.Equals(FieldName.TokenCode, null, "Smith"))), true },
+            { new SearchParameterExpression(NameParam, Expression.Equals(FieldName.TokenCode, null, "Smith")), false },
+        };
+
+        private static UnionExpression BuildBareUnion()
+        {
+            var endOfDay = new DateTimeOffset(2016, 7, 6, 23, 59, 59, TimeSpan.Zero);
+            var startOfDay = new DateTimeOffset(2016, 7, 6, 0, 0, 0, TimeSpan.Zero);
+            var shortBranch = new SearchParameterExpression(BirthdateParam, Expression.Equals(FieldName.DateTimeEnd, null, endOfDay));
+            var longBranch = new SearchParameterExpression(BirthdateParam, Expression.GreaterThanOrEqual(FieldName.DateTimeStart, null, startOfDay));
+            return Expression.Union(UnionOperator.All, new Expression[] { shortBranch, longBranch });
+        }
+
+        private static UnionExpression BuildSmartV2Union()
+        {
+            var branch = new SearchParameterExpression(NameParam, Expression.Equals(FieldName.TokenCode, null, "Smith"));
+            return Expression.Union(UnionOperator.All, new Expression[] { branch });
+        }
+
+        [Theory]
+        [MemberData(nameof(UnionDetectionCases))]
+        public void HasUnionAllExpression_DetectsBareAndNestedUnions(Expression predicate, bool expected)
+        {
+            var tableExpr = new SearchParamTableExpression(
+                DateTimeQueryGenerator.Instance,
+                predicate,
+                SearchParamTableExpressionKind.Normal);
+
+            Assert.Equal(expected, tableExpr.HasUnionAllExpression());
+        }
+
+        [Fact]
+        public void SplitExpressions_WhenPredicateContainsSingleUnionChild_ReturnsTrueWithUnionAndNullRemainder()
+        {
+            var union = BuildBareUnion();
+            var tableExpr = new SearchParamTableExpression(
+                DateTimeQueryGenerator.Instance,
+                Expression.And(union),
+                SearchParamTableExpressionKind.Normal);
+
+            bool result = tableExpr.SplitExpressions(out UnionExpression outUnion, out SearchParamTableExpression outRemainder);
+
+            Assert.True(result);
+            Assert.Same(union, outUnion);
+            Assert.Null(outRemainder);
+        }
+
+        [Fact]
+        public void SplitExpressions_WhenPredicateIsBareUnion_ReturnsTrueWithUnionAndNullRemainder()
+        {
+            var union = BuildBareUnion();
+            var tableExpr = new SearchParamTableExpression(
+                DateTimeQueryGenerator.Instance,
+                union,
+                SearchParamTableExpressionKind.Normal);
+
+            bool result = tableExpr.SplitExpressions(out UnionExpression outUnion, out SearchParamTableExpression outRemainder);
+
+            Assert.True(result);
+            Assert.Same(union, outUnion);
+            Assert.Null(outRemainder);
+        }
+
+        [Fact]
+        public void SplitExpressions_WhenUnionIsMarkedAsNoSplit_StillSeparatesUnion()
+        {
+            // DoNotSplitIntoSeparateCtes controls only how the union is lowered to SQL (a single inline UNION ALL CTE).
+            // It no longer prevents the union from being separated from its siblings, so the union is still returned.
+            var union = BuildBareUnion();
+            union.DoNotSplitIntoSeparateCtes = true;
+
+            var tableExpr = new SearchParamTableExpression(
+                DateTimeQueryGenerator.Instance,
+                Expression.And(union),
+                SearchParamTableExpressionKind.Normal);
+
+            bool result = tableExpr.SplitExpressions(out UnionExpression outUnion, out SearchParamTableExpression outRemainder);
+
+            Assert.True(result);
+            Assert.Same(union, outUnion);
+            Assert.True(outUnion.DoNotSplitIntoSeparateCtes);
+            Assert.Null(outRemainder);
+        }
+
+        [Fact]
+        public void SplitExpressions_WhenUnionNestedInMultiaryWithSiblings_ReturnsTrueWithUnionAndRemainder()
+        {
+            var union = BuildSmartV2Union();
+            var otherExpr = new SearchParameterExpression(NameParam, Expression.Equals(FieldName.TokenCode, null, "Smith"));
+            var multiary = Expression.And(union, otherExpr);
+
+            var tableExpr = new SearchParamTableExpression(
+                DateTimeQueryGenerator.Instance,
+                multiary,
+                SearchParamTableExpressionKind.Normal);
+
+            bool result = tableExpr.SplitExpressions(out UnionExpression outUnion, out SearchParamTableExpression outRemainder);
+
+            Assert.True(result);
+            Assert.Same(union, outUnion);
+            Assert.NotNull(outRemainder);
+            var remainderAnd = Assert.IsType<MultiaryExpression>(outRemainder.Predicate);
+            Assert.DoesNotContain(remainderAnd.Expressions, e => e is UnionExpression);
+        }
+
+        [Fact]
+        public void SplitExpressions_WhenPredicateIsPlainSearchParameter_ReturnsFalse()
+        {
+            var predicate = new SearchParameterExpression(NameParam, Expression.Equals(FieldName.TokenCode, null, "Smith"));
+            var tableExpr = new SearchParamTableExpression(
+                DateTimeQueryGenerator.Instance,
+                predicate,
+                SearchParamTableExpressionKind.Normal);
+
+            bool result = tableExpr.SplitExpressions(out UnionExpression outUnion, out SearchParamTableExpression outRemainder);
+
+            Assert.False(result);
+            Assert.Null(outUnion);
+            Assert.Null(outRemainder);
+        }
+
+        [Fact]
+        public void SortExpressionsByQueryLogic_WhenRegularUnionPrecedesConcatenationPair_ThenUnionMovesFirstAndPairStaysAdjacent()
+        {
+            var leadingNormal = BuildNormalTableExpression("leading");
+            var regularUnion = new SearchParamTableExpression(
+                DateTimeQueryGenerator.Instance,
+                Expression.And(BuildBareUnion()),
+                SearchParamTableExpressionKind.Normal);
+            var normalSibling = BuildNormalTableExpression("sibling-normal");
+            var concatenationSibling = new SearchParamTableExpression(
+                DateTimeQueryGenerator.Instance,
+                new SearchParameterExpression(NameParam, Expression.Equals(FieldName.TokenCode, null, "sibling-concat")),
+                SearchParamTableExpressionKind.Concatenation);
+            var trailingNormal = BuildNormalTableExpression("trailing");
+
+            var input = new List<SearchParamTableExpression>
+            {
+                leadingNormal,
+                regularUnion,
+                normalSibling,
+                concatenationSibling,
+                trailingNormal,
+            };
+
+            IReadOnlyList<SearchParamTableExpression> sorted = input.SortExpressionsByQueryLogic();
+
+            Assert.Same(regularUnion, sorted[0]);
+            Assert.Same(leadingNormal, sorted[1]);
+            Assert.Same(normalSibling, sorted[2]);
+            Assert.Same(concatenationSibling, sorted[3]);
+            Assert.Same(trailingNormal, sorted[4]);
+            Assert.Equal(SearchParamTableExpressionKind.Concatenation, sorted[3].Kind);
+        }
+
+        private static SearchParamTableExpression BuildNormalTableExpression(string code)
+        {
+            return new SearchParamTableExpression(
+                DateTimeQueryGenerator.Instance,
+                new SearchParameterExpression(NameParam, Expression.Equals(FieldName.TokenCode, null, code)),
+                SearchParamTableExpressionKind.Normal);
+        }
+    }
+}

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

@@ -35,6 +35,22 @@ namespace Microsoft.Health.Fhir.SqlServer.UnitTests.Features.Search;
 [Trait(Traits.Category, Categories.Search)]
 public class SqlQueryGeneratorTests
 {
+    private static readonly SearchParameterInfo BirthdateParam = new(
+        "birthdate",
+        "birthdate",
+        SearchParamType.Date,
+        new Uri("http://hl7.org/fhir/SearchParameter/individual-birthdate"),
+        expression: "Patient.birthDate",
+        baseResourceTypes: new[] { "Patient" });
+
+    private static readonly SearchParameterInfo NameParam = new(
+        "name",
+        "name",
+        SearchParamType.String,
+        new Uri("http://hl7.org/fhir/SearchParameter/individual-name"),
+        expression: "Patient.name",
+        baseResourceTypes: new[] { "Patient" });
+
     private readonly ISqlServerFhirModel _fhirModel;
     private readonly SearchParamTableExpressionQueryGeneratorFactory _queryGeneratorFactory;
     private readonly SchemaInformation _schemaInformation = new(SchemaVersionConstants.Min, SchemaVersionConstants.Max);
@@ -213,4 +229,64 @@ public void GivenReferenceSearchParameterWithMultipleTargetTypes_WhenSqlGenerate
         _fhirModel.Received(1).TryGetResourceTypeId("Patient", out Arg.Any<short>());
         _fhirModel.Received(1).TryGetResourceTypeId("Practitioner", out Arg.Any<short>());
     }
+
+    [Fact]
+    public void GivenNoSplitDayUnionWithSibling_WhenSqlGenerated_ThenUnionIsASingleCteWithInlineUnionAll()
+    {
+        UnionExpression union = BuildDaySplitUnion(doNotSplit: true);
+
+        SqlRootExpression sqlExpression = BuildUnionWithSibling(union);
+        SearchOptions searchOptions = new() { Sort = [], ResourceVersionTypes = ResourceVersionType.Latest };
+
+        _queryGenerator.VisitSqlRoot(sqlExpression, searchOptions);
+        string sql = _strBuilder.ToString();
+
+        // The union is lowered to a single CTE (cte0) whose body combines both branches with inline UNION ALL,
+        // and the sibling (cte1) intersects it. There is no separate per-branch or aggregating CTE, so cte2 never appears.
+        Assert.Contains("UNION ALL", sql);
+        Assert.Contains("cte0", sql);
+        Assert.Contains("cte1", sql);
+        Assert.DoesNotContain("cte2", sql);
+        Assert.Contains("EXISTS (SELECT * FROM cte0", sql);
+    }
+
+    [Fact]
+    public void GivenDayUnionWithoutNoSplitFlag_WhenSqlGenerated_ThenUnionFansOutIntoSeparateCtes()
+    {
+        UnionExpression union = BuildDaySplitUnion(doNotSplit: false);
+
+        SqlRootExpression sqlExpression = BuildUnionWithSibling(union);
+        SearchOptions searchOptions = new() { Sort = [], ResourceVersionTypes = ResourceVersionType.Latest };
+
+        _queryGenerator.VisitSqlRoot(sqlExpression, searchOptions);
+        string sql = _strBuilder.ToString();
+
+        // Without the flag the union fans out into one CTE per branch (cte0, cte1) plus an aggregating CTE (cte2)
+        // that UNION ALLs them, followed by the sibling CTE (cte3).
+        Assert.Contains("SELECT * FROM cte0", sql);
+        Assert.Contains("cte3", sql);
+    }
+
+    private static UnionExpression BuildDaySplitUnion(bool doNotSplit)
+    {
+        var endOfDay = new DateTimeOffset(2016, 7, 6, 23, 59, 59, TimeSpan.Zero);
+        var startOfDay = new DateTimeOffset(2016, 7, 6, 0, 0, 0, TimeSpan.Zero);
+        var shortBranch = new SearchParameterExpression(BirthdateParam, Expression.Equals(FieldName.DateTimeEnd, null, endOfDay));
+        var longBranch = new SearchParameterExpression(BirthdateParam, Expression.GreaterThanOrEqual(FieldName.DateTimeStart, null, startOfDay));
+        var union = Expression.Union(UnionOperator.All, new Expression[] { shortBranch, longBranch });
+        union.DoNotSplitIntoSeparateCtes = doNotSplit;
+        return union;
+    }
+
+    private static SqlRootExpression BuildUnionWithSibling(UnionExpression union)
+    {
+        // The runtime shape: the union is the entire predicate of its own table expression (a bare UnionExpression).
+        var unionTable = new SearchParamTableExpression(DateTimeQueryGenerator.Instance, union, SearchParamTableExpressionKind.Normal);
+        var siblingTable = new SearchParamTableExpression(
+            StringQueryGenerator.Instance,
+            new SearchParameterExpression(NameParam, Expression.StringEquals(FieldName.String, null, "Smith", false)),
+            SearchParamTableExpressionKind.Normal);
+
+        return new SqlRootExpression(new[] { unionTable, siblingTable }, new List<SearchParameterExpressionBase>());
+    }
 }

src/Microsoft.Health.Fhir.SqlServer/Features/Search/Expressions/SearchParamTableExpressionExtensions.cs

@@ -20,6 +20,13 @@ internal static class SearchParamTableExpressionExtensions
         /// <param name="expression">Instance of <see cref="SearchParamTableExpression"/> under evaluation.</param>
         public static bool HasUnionAllExpression(this SearchParamTableExpression expression)
         {
+            // The predicate may either BE a UnionExpression (e.g. a scalar date day-split that replaced a single
+            // SearchParameterExpression) or be a container (e.g. And(union, ...)) that holds one.
+            if (expression.Predicate is UnionExpression)
+            {
+                return true;
+            }
+
             IExpressionsContainer expressionContainer = expression.Predicate as IExpressionsContainer;
             return expressionContainer?.Expressions.Any(e => e is UnionExpression) ?? false;
         }
@@ -36,6 +43,14 @@ public static bool SplitExpressions(this SearchParamTableExpression expression,
             unionExpression = null;
             allOtherRemainingExpressions = null;
 
+            // The union may be the entire predicate (e.g. a scalar date day-split that replaced a single
+            // SearchParameterExpression). In that case there are no sibling expressions to separate out.
+            if (expression.Predicate is UnionExpression bareUnionExpression)
+            {
+                unionExpression = bareUnionExpression;
+                return true;
+            }
+
             IExpressionsContainer expressionContainer = expression.Predicate as IExpressionsContainer;
 
             if (expressionContainer != null)