Fix continuation token issues with chained searches (#1106)

* ExportDataValidationTests fail if there are duplicate IDs

* Promote denormalized predicates to chained table expressions

* Added some comments

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

@@ -18,7 +18,7 @@ namespace Microsoft.Health.Fhir.SqlServer.UnitTests.Features.Search.Expressions
     public class DenormalizedPredicateRewriterTests
     {
         [Fact]
-        public void GivenExpressionWithNoTableExpressions_WhenRewritted_ReturnsOriginalExpression()
+        public void GivenExpressionWithNoTableExpressions_WhenRewritten_ReturnsOriginalExpression()
         {
             var inputExpression = SqlRootExpression.WithDenormalizedExpressions(
                 Expression.Equals(FieldName.Number, null, 1));
@@ -27,7 +27,7 @@ public void GivenExpressionWithNoTableExpressions_WhenRewritted_ReturnsOriginalE
         }
 
         [Fact]
-        public void GivenExpressionWithNoDenormalizedExpressions_WhenRewritted_ReturnsOriginalExpression()
+        public void GivenExpressionWithNoDenormalizedExpressions_WhenRewritten_ReturnsOriginalExpression()
         {
             var inputExpression = SqlRootExpression.WithTableExpressions(
                 new TableExpression(null, null, null, TableExpressionKind.Normal));
@@ -104,5 +104,24 @@ public void GivenExpressionWithMultipleDenormalizedExpressions_WhenRewritten_Den
             var visitedExpression = (SqlRootExpression)inputExpression.AcceptVisitor(DenormalizedPredicateRewriter.Instance);
             Assert.Equal(inputExpression.DenormalizedExpressions[1], visitedExpression.DenormalizedExpressions[0]);
         }
+
+        [Theory]
+        [InlineData(SearchParameterNames.ResourceType)]
+        [InlineData(SqlSearchParameters.ResourceSurrogateIdParameterName)]
+        public void GivenExpressionWithDenormalizedAndChainedExpressions_WhenRewritten_DenormalisedPredicatesPromotedToChainTableExpression(string paramName)
+        {
+            var inputExpression = new SqlRootExpression(
+                new List<TableExpression>
+                {
+                    new TableExpression(null, null, null, TableExpressionKind.Chain, 1),
+                },
+                new List<Expression>
+                {
+                    new SearchParameterExpression(new SearchParameterInfo(paramName), Expression.Equals(FieldName.String, null, "ExtractableTestParamValue")),
+                });
+            var visitedExpression = (SqlRootExpression)inputExpression.AcceptVisitor(DenormalizedPredicateRewriter.Instance);
+            Assert.Equal(inputExpression.DenormalizedExpressions[0], visitedExpression.TableExpressions[0].DenormalizedPredicateOnChainRoot);
+            Assert.Equal(inputExpression.TableExpressions[0].ChainLevel, visitedExpression.TableExpressions[0].ChainLevel);
+        }
     }
 }

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

@@ -15,12 +15,23 @@ namespace Microsoft.Health.Fhir.SqlServer.Features.Search.Expressions
     /// </summary>
     internal class TableExpression : Expression
     {
+        /// <summary>
+        /// Creates a new instance of the <see cref="TableExpression"/> class.
+        /// </summary>
+        /// <param name="searchParameterQueryGenerator">The search parameter query generator</param>
+        /// <param name="normalizedPredicate">The search expression over a columns belonging exclusively to a search parameter table.
+        /// Applies to the chain target if a chained expression.</param>
+        /// <param name="denormalizedPredicate">The search expression over columns that exist the Resource table. Applies to the chain target if a chained expression.</param>
+        /// <param name="kind">The table expression kind.</param>
+        /// <param name="chainLevel">The nesting chain nesting level of the current expression. 0 if not a chain expression.</param>
+        /// <param name="denormalizedPredicateOnChainRoot">The search expression over columns that exist the Resource table. Applies to the chain root if a chained expression.</param>
         public TableExpression(
             NormalizedSearchParameterQueryGenerator searchParameterQueryGenerator,
             Expression normalizedPredicate,
             Expression denormalizedPredicate,
             TableExpressionKind kind,
-            int chainLevel = 0)
+            int chainLevel = 0,
+            Expression denormalizedPredicateOnChainRoot = null)
         {
             switch (normalizedPredicate)
             {
@@ -39,18 +50,34 @@ internal class TableExpression : Expression
             DenormalizedPredicate = denormalizedPredicate;
             Kind = kind;
             ChainLevel = chainLevel;
+            DenormalizedPredicateOnChainRoot = denormalizedPredicateOnChainRoot;
         }
 
         public TableExpressionKind Kind { get; }
 
+        /// <summary>
+        /// The nesting chain nesting level of the current expression. 0 if not a chain expression.
+        /// </summary>
         public int ChainLevel { get; }
 
         public NormalizedSearchParameterQueryGenerator SearchParameterQueryGenerator { get; }
 
+        /// <summary>
+        /// The search expression over a columns belonging exclusively to a search parameter table.
+        /// Applies to the chain target if a chained expression.
+        /// </summary>
         public Expression NormalizedPredicate { get; }
 
+        /// <summary>
+        /// The search expression over columns that exist the Resource table. Applies to the chain target if a chained expression.
+        /// </summary>
         public Expression DenormalizedPredicate { get; }
 
+        /// <summary>
+        /// The search expression over columns that exist the Resource table. Applies to the chain root if a chained expression.
+        /// </summary>
+        public Expression DenormalizedPredicateOnChainRoot { get; }
+
         public override TOutput AcceptVisitor<TContext, TOutput>(IExpressionVisitor<TContext, TOutput> visitor, TContext context)
         {
             return AcceptVisitor((ISqlExpressionVisitor<TContext, TOutput>)visitor, context);
@@ -63,7 +90,7 @@ internal class TableExpression : Expression
 
         public override string ToString()
         {
-            return $"(Table {Kind} {(ChainLevel == 0 ? null : $"ChainLevel:{ChainLevel} ")}{SearchParameterQueryGenerator?.Table} Normalized:{NormalizedPredicate} Denormalized:{DenormalizedPredicate})";
+            return $"(Table {Kind} {(ChainLevel == 0 ? null : $"ChainLevel:{ChainLevel} ")}{SearchParameterQueryGenerator?.Table} Normalized:{NormalizedPredicate} Denormalized:{DenormalizedPredicate} DenormalizedOnChainRoot:{DenormalizedPredicateOnChainRoot})";
         }
     }
 }

src/Microsoft.Health.Fhir.SqlServer/Features/Search/Expressions/Visitors/ChainFlatteningRewriter.cs

@@ -66,7 +66,8 @@ public override Expression VisitChained(ChainedExpression expression, (TableExpr
                     expression,
                     denormalizedPredicate: normalizedParameterQueryGenerator == null ? expression.Expression : null,
                     TableExpressionKind.Chain,
-                    context.chainLevel);
+                    context.chainLevel,
+                    thisTableExpression?.DenormalizedPredicateOnChainRoot);
             }
 
             if (normalizedParameterQueryGenerator == null)

src/Microsoft.Health.Fhir.SqlServer/Features/Search/Expressions/Visitors/DenormalizedPredicateRewriter.cs

@@ -22,7 +22,7 @@ internal class DenormalizedPredicateRewriter : ExpressionRewriterWithInitialCont
 
         public Expression VisitSqlRoot(SqlRootExpression expression, object context)
         {
-            if (expression.TableExpressions.Count == 0 || expression.DenormalizedExpressions.Count == 0 || expression.TableExpressions.All(t => t.Kind == TableExpressionKind.Chain) || expression.TableExpressions.All(t => t.Kind == TableExpressionKind.Include))
+            if (expression.TableExpressions.Count == 0 || expression.DenormalizedExpressions.Count == 0 || expression.TableExpressions.All(t => t.Kind == TableExpressionKind.Include))
             {
                 return expression;
             }
@@ -59,10 +59,14 @@ public Expression VisitSqlRoot(SqlRootExpression expression, object context)
             var newTableExpressions = new List<TableExpression>(expression.TableExpressions.Count);
             foreach (var tableExpression in expression.TableExpressions)
             {
-                if (tableExpression.Kind == TableExpressionKind.Chain || tableExpression.Kind == TableExpressionKind.Include)
+                if (tableExpression.Kind == TableExpressionKind.Include)
                 {
                     newTableExpressions.Add(tableExpression);
                 }
+                else if (tableExpression.Kind == TableExpressionKind.Chain)
+                {
+                    newTableExpressions.Add(new TableExpression(tableExpression.SearchParameterQueryGenerator, tableExpression.NormalizedPredicate, tableExpression.DenormalizedPredicate, tableExpression.Kind, chainLevel: tableExpression.ChainLevel, denormalizedPredicateOnChainRoot: extractedDenormalizedExpression));
+                }
                 else
                 {
                     Expression newDenormalizedPredicate = tableExpression.DenormalizedPredicate == null