Ensure prefix length does not exceed configured maximum (facebookincu…

…bator#11496)

Summary:
The check can't make sure the normalizedKeySize is less than or equal to
maxNormalizedKeySize. That is possible
`normalizedKeySize + encodedSize.value() > maxNormalizedKeySize` but related
column is included in the prefix.
Move the check to fix this.
This is found when doing the PR: facebookincubator#11272

Pull Request resolved: facebookincubator#11496

Reviewed By: tanjialiang

Differential Revision: D65850949

Pulled By: xiaoxmeng

fbshipit-source-id: 2dc5fcb93f0960544b9df3f2471caa2923c70975

velox/exec/PrefixSort.cpp

@@ -136,12 +136,10 @@ PrefixSortLayout PrefixSortLayout::makeSortLayout(
   uint32_t normalizedKeySize{0};
   uint32_t numNormalizedKeys{0};
   for (auto i = 0; i < numKeys; ++i) {
-    if (normalizedKeySize > maxNormalizedKeySize) {
-      break;
-    }
     const std::optional<uint32_t> encodedSize =
         PrefixSortEncoder::encodedSize(types[i]->kind());
-    if (!encodedSize.has_value()) {
+    if (!encodedSize.has_value() ||
+        normalizedKeySize + encodedSize.value() > maxNormalizedKeySize) {
       break;
     }
     prefixOffsets.push_back(normalizedKeySize);

velox/exec/tests/PrefixSortTest.cpp

@@ -280,5 +280,36 @@ TEST_F(PrefixSortTest, fuzzMulti) {
     testPrefixSort({kDesc, kDesc}, data);
   }
 }
+
+TEST_F(PrefixSortTest, checkMaxNormalizedKeySizeForMultipleKeys) {
+  // Test the normalizedKeySize doesn't exceed the MaxNormalizedKeySize.
+  // The normalizedKeySize for BIGINT should be 8 + 1.
+  std::vector<TypePtr> keyTypes = {BIGINT(), BIGINT()};
+  std::vector<CompareFlags> compareFlags = {kAsc, kDesc};
+  auto sortLayout = PrefixSortLayout::makeSortLayout(keyTypes, compareFlags, 8);
+  ASSERT_FALSE(sortLayout.hasNormalizedKeys);
+
+  auto sortLayoutOneKey =
+      PrefixSortLayout::makeSortLayout(keyTypes, compareFlags, 9);
+  ASSERT_TRUE(sortLayoutOneKey.hasNormalizedKeys);
+  ASSERT_TRUE(sortLayoutOneKey.hasNonNormalizedKey);
+  ASSERT_EQ(sortLayoutOneKey.prefixOffsets.size(), 1);
+  ASSERT_EQ(sortLayoutOneKey.prefixOffsets[0], 0);
+
+  auto sortLayoutOneKey1 =
+      PrefixSortLayout::makeSortLayout(keyTypes, compareFlags, 17);
+  ASSERT_TRUE(sortLayoutOneKey1.hasNormalizedKeys);
+  ASSERT_TRUE(sortLayoutOneKey1.hasNonNormalizedKey);
+  ASSERT_EQ(sortLayoutOneKey1.prefixOffsets.size(), 1);
+  ASSERT_EQ(sortLayoutOneKey1.prefixOffsets[0], 0);
+
+  auto sortLayoutTwoKeys =
+      PrefixSortLayout::makeSortLayout(keyTypes, compareFlags, 18);
+  ASSERT_TRUE(sortLayoutTwoKeys.hasNormalizedKeys);
+  ASSERT_FALSE(sortLayoutTwoKeys.hasNonNormalizedKey);
+  ASSERT_EQ(sortLayoutTwoKeys.prefixOffsets.size(), 2);
+  ASSERT_EQ(sortLayoutTwoKeys.prefixOffsets[0], 0);
+  ASSERT_EQ(sortLayoutTwoKeys.prefixOffsets[1], 9);
+}
 } // namespace
 } // namespace facebook::velox::exec::prefixsort::test