diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/CosmosExceptionExtensions.cs b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/CosmosExceptionExtensions.cs
index 98e9023226..dca1ef488d 100644
--- a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/CosmosExceptionExtensions.cs
+++ b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/CosmosExceptionExtensions.cs
@@ -9,10 +9,20 @@ namespace Microsoft.Azure.Cosmos.Query.Core.Pipeline
     internal static class CosmosExceptionExtensions
     {
         public static bool IsPartitionSplitException(this Exception ex)
+        {
+            if (ex != null)
+            {
+                return IsPartitionSplitException(ex as CosmosException);
+            }
+
+            return false;
+        }
+
+        public static bool IsPartitionSplitException(this CosmosException ex)
         {
             return ex is CosmosException cosmosException
-                && (cosmosException.StatusCode == System.Net.HttpStatusCode.Gone)
-                && (cosmosException.SubStatusCode == (int)Documents.SubStatusCodes.PartitionKeyRangeGone);
+            && (cosmosException.StatusCode == System.Net.HttpStatusCode.Gone)
+            && (cosmosException.SubStatusCode == (int)Documents.SubStatusCodes.PartitionKeyRangeGone);
         }
     }
 }
diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/CosmosQueryExecutionContextFactory.cs b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/CosmosQueryExecutionContextFactory.cs
index 82cde7df15..2476d9d701 100644
--- a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/CosmosQueryExecutionContextFactory.cs
+++ b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/CosmosQueryExecutionContextFactory.cs
@@ -305,19 +305,22 @@ private static async Task<TryCatch<IQueryPipelineStage>> TryCreateFromPartitione
                     trace,
                     cancellationToken);
             }
-            else if (createPassthroughQuery)
-            {
-                SetTestInjectionPipelineType(inputParameters, Passthrough);
-
-                tryCreatePipelineStage = CosmosQueryExecutionContextFactory.TryCreatePassthroughQueryExecutionContext(
-                    documentContainer,
-                    inputParameters,
-                    targetRanges,
-                    cancellationToken);
-            }
             else
             {
-                tryCreatePipelineStage = TryCreateSpecializedDocumentQueryExecutionContext(documentContainer, cosmosQueryContext, inputParameters, targetRanges, partitionedQueryExecutionInfo, cancellationToken);
+                if (createPassthroughQuery)
+                {
+                    SetTestInjectionPipelineType(inputParameters, Passthrough);
+
+                    tryCreatePipelineStage = CosmosQueryExecutionContextFactory.TryCreatePassthroughQueryExecutionContext(
+                        documentContainer,
+                        inputParameters,
+                        targetRanges,
+                        cancellationToken);
+                }
+                else
+                {
+                    tryCreatePipelineStage = TryCreateSpecializedDocumentQueryExecutionContext(documentContainer, cosmosQueryContext, inputParameters, targetRanges, partitionedQueryExecutionInfo, cancellationToken);
+                }
             }
 
             return tryCreatePipelineStage;
diff --git a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Query/OptimisticDirectExecutionQueryBaselineTests.cs b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Query/OptimisticDirectExecutionQueryBaselineTests.cs
index c0644045e5..6737a3d338 100644
--- a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Query/OptimisticDirectExecutionQueryBaselineTests.cs
+++ b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Query/OptimisticDirectExecutionQueryBaselineTests.cs
@@ -623,7 +623,8 @@ public async Task<Exception> ShouldReturnFailure()
                         activityId: "0f8fad5b-d9cb-469f-a165-70867728950e",
                         requestCharge: default);
                 }
-                else if (this.IsFailedFallbackPipelineTest && this.GoneExceptionCreated && !this.TooManyRequestsFailureCreated)
+                
+                if (this.IsFailedFallbackPipelineTest && this.GoneExceptionCreated && !this.TooManyRequestsFailureCreated)
                 {
                     this.TooManyRequestsFailureCreated = true;
                     return new CosmosException(
@@ -633,10 +634,8 @@ public async Task<Exception> ShouldReturnFailure()
                             activityId: "111fad5b-d9cb-469f-a165-70867728950e",
                             requestCharge: 0);
                 }
-                else
-                {
-                    return null;
-                }
+
+                return null;
             }
         }
     }