feat(prefix sort): Support string type key in prefix sort (#11527)

Summary: Support string type key in PrefixSort. We introduce the configuration parameter `prefixsort_max_string_length`, which sets the maximum prefix length for strings. The implementation dynamically determines the prefix length by comparing the configured maximum with the actual maximum column length from RowContainer, using the smaller of the two. This ensures efficient and flexible prefix sorting for variable-length string types. Default value of `prefixsort_max_string_length` is 16. Perf result: ``` StdSort_no-payloads_1_varchar_1k 158.23ns 6.32M PrefixSort 415.89% 38.05ns 26.28M StdSort_no-payloads_2_varchar_1k 186.19ns 5.37M PrefixSort 197.21% 94.41ns 10.59M StdSort_no-payloads_3_varchar_1k 197.90ns 5.05M PrefixSort 148.93% 132.89ns 7.53M StdSort_no-payloads_4_varchar_1k 211.35ns 4.73M PrefixSort 126.75% 166.75ns 6.00M StdSort_no-payloads_1_varchar_10k 257.23ns 3.89M PrefixSort 358.08% 71.84ns 13.92M StdSort_no-payloads_2_varchar_10k 272.61ns 3.67M PrefixSort 227.46% 119.85ns 8.34M StdSort_no-payloads_3_varchar_10k 295.37ns 3.39M PrefixSort 170.20% 173.55ns 5.76M StdSort_no-payloads_4_varchar_10k 319.42ns 3.13M PrefixSort 152.60% 209.31ns 4.78M StdSort_no-payloads_1_varchar_100k 348.19ns 2.87M PrefixSort 403.18% 86.36ns 11.58M StdSort_no-payloads_2_varchar_100k 409.94ns 2.44M PrefixSort 261.56% 156.73ns 6.38M StdSort_no-payloads_3_varchar_100k 469.93ns 2.13M PrefixSort 206.32% 227.76ns 4.39M StdSort_no-payloads_4_varchar_100k 526.94ns 1.90M PrefixSort 186.66% 282.29ns 3.54M StdSort_no-payloads_1_varchar_1000k 780.28ns 1.28M PrefixSort 627.88% 124.27ns 8.05M StdSort_no-payloads_2_varchar_1000k 976.32ns 1.02M PrefixSort 491.56% 198.62ns 5.03M StdSort_no-payloads_3_varchar_1000k 1.08us 928.51K PrefixSort 376.44% 286.10ns 3.50M StdSort_no-payloads_4_varchar_1000k 1.12us 889.85K PrefixSort 321.50% 349.54ns 2.86M ``` Pull Request resolved: #11527 Reviewed By: Yuhta Differential Revision: D67149095 Pulled By: xiaoxmeng fbshipit-source-id: 79f02c81165a873aa8068260b5580850f30a4fc5
facebookincubator · Dec 13, 2024 · 960c2af · 960c2af
1 parent e86ff05
commit 960c2af
Show file tree

Hide file tree

Showing 19 changed files with 505 additions and 154 deletions.
diff --git a/velox/common/base/PrefixSortConfig.h b/velox/common/base/PrefixSortConfig.h
@@ -24,9 +24,13 @@ namespace facebook::velox::common {
 struct PrefixSortConfig {
   PrefixSortConfig() = default;
 
-  PrefixSortConfig(uint32_t _maxNormalizedKeyBytes, uint32_t _minNumRows)
+  PrefixSortConfig(
+      uint32_t _maxNormalizedKeyBytes,
+      uint32_t _minNumRows,
+      uint32_t _maxStringPrefixLength)
       : maxNormalizedKeyBytes(_maxNormalizedKeyBytes),
-        minNumRows(_minNumRows) {}
+        minNumRows(_minNumRows),
+        maxStringPrefixLength(_maxStringPrefixLength) {}
 
   /// Maximum bytes that can be used to store normalized keys in prefix-sort
   /// buffer per entry. Same with QueryConfig kPrefixSortNormalizedKeyMaxBytes.
@@ -35,5 +39,9 @@ struct PrefixSortConfig {
   /// Minimum number of rows to apply prefix sort. Prefix sort does not perform
   /// with small datasets.
   uint32_t minNumRows{128};
+
+  /// Maximum number of bytes to be stored in prefix-sort buffer for a string
+  /// column.
+  uint32_t maxStringPrefixLength{16};
 };
 } // namespace facebook::velox::common
diff --git a/velox/core/QueryConfig.h b/velox/core/QueryConfig.h
@@ -378,6 +378,11 @@ class QueryConfig {
   /// derived using micro-benchmarking.
   static constexpr const char* kPrefixSortMinRows = "prefixsort_min_rows";
 
+  /// Maximum number of bytes to be stored in prefix-sort buffer for a string
+  /// key.
+  static constexpr const char* kPrefixSortMaxStringPrefixLength =
+      "prefixsort_max_string_prefix_length";
+
   /// Enable query tracing flag.
   static constexpr const char* kQueryTraceEnabled = "query_trace_enabled";
 
@@ -844,6 +849,10 @@ class QueryConfig {
     return get<uint32_t>(kPrefixSortMinRows, 128);
   }
 
+  uint32_t prefixSortMaxStringPrefixLength() const {
+    return get<uint32_t>(kPrefixSortMaxStringPrefixLength, 16);
+  }
+
   double scaleWriterRebalanceMaxMemoryUsageRatio() const {
     return get<double>(kScaleWriterRebalanceMaxMemoryUsageRatio, 0.7);
   }

diff --git a/velox/docs/configs.rst b/velox/docs/configs.rst
@@ -144,6 +144,10 @@ Generic Configuration
      - integer
      - 128
      - Minimum number of rows to use prefix-sort. The default value has been derived using micro-benchmarking.
+   * - prefixsort_max_string_prefix_length
+     - integer
+     - 16
+     - Byte length of the string prefix stored in the prefix-sort buffer. This doesn't include the null byte.
 
 .. _expression-evaluation-conf:
 

diff --git a/velox/exec/Driver.h b/velox/exec/Driver.h
@@ -300,7 +300,8 @@ struct DriverCtx {
   common::PrefixSortConfig prefixSortConfig() const {
     return common::PrefixSortConfig{
         queryConfig().prefixSortNormalizedKeyMaxBytes(),
-        queryConfig().prefixSortMinRows()};
+        queryConfig().prefixSortMinRows(),
+        queryConfig().prefixSortMaxStringPrefixLength()};
   }
 };
 

diff --git a/velox/exec/PrefixSort.cpp b/velox/exec/PrefixSort.cpp
@@ -40,7 +40,9 @@ FOLLY_ALWAYS_INLINE void encodeRowColumn(
     value = *(reinterpret_cast<T*>(row + rowColumn.offset()));
   }
   prefixSortLayout.encoders[index].encode(
-      value, prefixBuffer + prefixSortLayout.prefixOffsets[index]);
+      value,
+      prefixBuffer + prefixSortLayout.prefixOffsets[index],
+      prefixSortLayout.encodeSizes[index]);
 }
 
 FOLLY_ALWAYS_INLINE void extractRowColumnToPrefix(
@@ -86,6 +88,13 @@ FOLLY_ALWAYS_INLINE void extractRowColumnToPrefix(
           prefixSortLayout, index, rowColumn, row, prefixBuffer);
       return;
     }
+    case TypeKind::VARCHAR:
+      [[fallthrough]];
+    case TypeKind::VARBINARY: {
+      encodeRowColumn<StringView>(
+          prefixSortLayout, index, rowColumn, row, prefixBuffer);
+      return;
+    }
     default:
       VELOX_UNSUPPORTED(
           "prefix-sort does not support type kind: {}",
@@ -127,31 +136,49 @@ compareByWord(uint64_t* left, uint64_t* right, int32_t bytes) {
 } // namespace
 
 // static.
-PrefixSortLayout PrefixSortLayout::makeSortLayout(
+PrefixSortLayout PrefixSortLayout::generate(
     const std::vector<TypePtr>& types,
     const std::vector<CompareFlags>& compareFlags,
-    uint32_t maxNormalizedKeySize) {
+    uint32_t maxNormalizedKeySize,
+    uint32_t maxStringPrefixLength,
+    const std::vector<std::optional<uint32_t>>& maxStringLengths) {
   const uint32_t numKeys = types.size();
   std::vector<uint32_t> prefixOffsets;
   prefixOffsets.reserve(numKeys);
+  std::vector<uint32_t> encodeSizes;
+  encodeSizes.reserve(numKeys);
   std::vector<PrefixSortEncoder> encoders;
   encoders.reserve(numKeys);
 
   // Calculate encoders and prefix-offsets, and stop the loop if a key that
-  // cannot be normalized is encountered.
+  // cannot be normalized is encountered or only partial data of a key is
+  // normalized.
   uint32_t normalizedKeySize{0};
   uint32_t numNormalizedKeys{0};
+
+  bool lastPrefixKeyPartial{false};
   for (auto i = 0; i < numKeys; ++i) {
-    const std::optional<uint32_t> encodedSize =
-        PrefixSortEncoder::encodedSize(types[i]->kind());
+    const std::optional<uint32_t> encodedSize = PrefixSortEncoder::encodedSize(
+        types[i]->kind(),
+        maxStringLengths[i].has_value()
+            ? std::min(maxStringLengths[i].value(), maxStringPrefixLength)
+            : maxStringPrefixLength);
     if (!encodedSize.has_value() ||
         normalizedKeySize + encodedSize.value() > maxNormalizedKeySize) {
       break;
     }
     prefixOffsets.push_back(normalizedKeySize);
     encoders.push_back({compareFlags[i].ascending, compareFlags[i].nullsFirst});
+    encodeSizes.push_back(encodedSize.value());
     normalizedKeySize += encodedSize.value();
     ++numNormalizedKeys;
+    if ((types[i]->kind() == TypeKind::VARCHAR ||
+         types[i]->kind() == TypeKind::VARBINARY) &&
+        (!maxStringLengths[i].has_value() ||
+         maxStringPrefixLength < maxStringLengths[i].value())) {
+      lastPrefixKeyPartial = true;
+      break;
+    }
   }
 
   const auto numPaddingBytes = alignmentPadding(normalizedKeySize, kAlignment);
@@ -165,7 +192,10 @@ PrefixSortLayout PrefixSortLayout::makeSortLayout(
       compareFlags,
       numNormalizedKeys != 0,
       numNormalizedKeys < numKeys,
+      /*nonPrefixSortStartIndex=*/
+      lastPrefixKeyPartial ? numNormalizedKeys - 1 : numNormalizedKeys,
       std::move(prefixOffsets),
+      std::move(encodeSizes),
       std::move(encoders),
       numPaddingBytes};
 }
@@ -177,8 +207,10 @@ void PrefixSortLayout::optimizeSortKeysOrder(
   std::vector<std::optional<uint32_t>> encodedKeySizes(
       rowType->size(), std::nullopt);
   for (const auto& projection : keyColumnProjections) {
+    // Set maxStringPrefixLength to UINT_MAX - 1 to ensure VARCHAR columns are
+    // placed after all other supported types and before un-supported types.
     encodedKeySizes[projection.inputChannel] = PrefixSortEncoder::encodedSize(
-        rowType->childAt(projection.inputChannel)->kind());
+        rowType->childAt(projection.inputChannel)->kind(), UINT_MAX - 1);
   }
 
   std::sort(
@@ -222,7 +254,8 @@ int PrefixSort::comparePartNormalizedKeys(char* left, char* right) {
   // If prefixes are equal, compare the remaining sort keys with rowContainer.
   char* leftRow = getRowAddrFromPrefixBuffer(left);
   char* rightRow = getRowAddrFromPrefixBuffer(right);
-  for (auto i = sortLayout_.numNormalizedKeys; i < sortLayout_.numKeys; ++i) {
+  for (auto i = sortLayout_.nonPrefixSortStartIndex; i < sortLayout_.numKeys;
+       ++i) {
     result = rowContainer_->compare(
         leftRow, rightRow, i, sortLayout_.compareFlags[i]);
     if (result != 0) {
@@ -276,9 +309,8 @@ uint32_t PrefixSort::maxRequiredBytes(
   if (rowContainer->numRows() < config.minNumRows) {
     return 0;
   }
-  VELOX_CHECK_EQ(rowContainer->keyTypes().size(), compareFlags.size());
-  const auto sortLayout = PrefixSortLayout::makeSortLayout(
-      rowContainer->keyTypes(), compareFlags, config.maxNormalizedKeyBytes);
+  const auto sortLayout =
+      generateSortLayout(rowContainer, compareFlags, config);
   if (!sortLayout.hasNormalizedKeys) {
     return 0;
   }
@@ -346,7 +378,8 @@ void PrefixSort::sortInternal(
           RuntimeCounter(
               sortLayout_.numNormalizedKeys, RuntimeCounter::Unit::kNone));
     }
-    if (sortLayout_.hasNonNormalizedKey) {
+    if (sortLayout_.hasNonNormalizedKey ||
+        sortLayout_.nonPrefixSortStartIndex < sortLayout_.numNormalizedKeys) {
       sortRunner.quickSort(
           prefixBufferStart, prefixBufferEnd, [&](char* lhs, char* rhs) {
             return comparePartNormalizedKeys(lhs, rhs);

diff --git a/velox/exec/PrefixSort.h b/velox/exec/PrefixSort.h
@@ -15,6 +15,8 @@
  */
 #pragma once
 
+#include <optional>
+
 #include "velox/common/base/PrefixSortConfig.h"
 #include "velox/exec/Operator.h"
 #include "velox/exec/RowContainer.h"
@@ -53,21 +55,31 @@ struct PrefixSortLayout {
   /// Whether the sort keys contains non-normalized key.
   const bool hasNonNormalizedKey;
 
+  /// Indicates the starting index for key comparison.
+  /// If the last key is only partially encoded in the prefix, start from
+  /// numNormalizedKeys - 1. Otherwise, start from numNormalizedKeys.
+  const uint32_t nonPrefixSortStartIndex;
+
   /// Offsets of normalized keys, used to find write locations when
   /// extracting columns
   const std::vector<uint32_t> prefixOffsets;
 
+  /// Sizes of normalized keys.
+  const std::vector<uint32_t> encodeSizes;
+
   /// The encoders for normalized keys.
   const std::vector<prefixsort::PrefixSortEncoder> encoders;
 
   /// The number of padding bytes to align each prefix encoded row size to 8
   /// for fast long compare.
   const int32_t numPaddingBytes;
 
-  static PrefixSortLayout makeSortLayout(
+  static PrefixSortLayout generate(
       const std::vector<TypePtr>& types,
       const std::vector<CompareFlags>& compareFlags,
-      uint32_t maxNormalizedKeySize);
+      uint32_t maxNormalizedKeySize,
+      uint32_t maxStringPrefixLength,
+      const std::vector<std::optional<uint32_t>>& maxStringLengths);
 
   /// Optimizes the order of sort key columns to maximize the number of prefix
   /// sort keys for acceleration. This only applies for use case which doesn't
@@ -121,10 +133,8 @@ class PrefixSort {
       stdSort(rows, rowContainer, compareFlags);
       return;
     }
-
-    VELOX_CHECK_EQ(rowContainer->keyTypes().size(), compareFlags.size());
-    const auto sortLayout = PrefixSortLayout::makeSortLayout(
-        rowContainer->keyTypes(), compareFlags, config.maxNormalizedKeyBytes);
+    const auto sortLayout =
+        generateSortLayout(rowContainer, compareFlags, config);
     // All keys can not normalize, skip the binary string compare opt.
     // Putting this outside sort-internal helps with stdSort.
     if (!sortLayout.hasNormalizedKeys) {
@@ -158,6 +168,33 @@ class PrefixSort {
       const RowContainer* rowContainer,
       const std::vector<CompareFlags>& compareFlags);
 
+  FOLLY_ALWAYS_INLINE static PrefixSortLayout generateSortLayout(
+      const RowContainer* rowContainer,
+      const std::vector<CompareFlags>& compareFlags,
+      const velox::common::PrefixSortConfig& config) {
+    const auto keyTypes = rowContainer->keyTypes();
+    VELOX_CHECK_EQ(keyTypes.size(), compareFlags.size());
+    std::vector<std::optional<uint32_t>> maxStringLengths;
+    maxStringLengths.reserve(keyTypes.size());
+    for (int i = 0; i < keyTypes.size(); ++i) {
+      std::optional<uint32_t> maxStringLength = std::nullopt;
+      if (keyTypes[i]->kind() == TypeKind::VARBINARY ||
+          keyTypes[i]->kind() == TypeKind::VARCHAR) {
+        const auto stats = rowContainer->columnStats(i);
+        if (stats.has_value()) {
+          maxStringLength = stats.value().maxBytes();
+        }
+      }
+      maxStringLengths.emplace_back(maxStringLength);
+    }
+    return PrefixSortLayout::generate(
+        keyTypes,
+        compareFlags,
+        config.maxNormalizedKeyBytes,
+        config.maxStringPrefixLength,
+        maxStringLengths);
+  }
+
   // Estimates the memory required for prefix sort such as prefix buffer and
   // swap buffer.
   uint32_t maxRequiredBytes() const;

diff --git a/velox/exec/RowContainer.cpp b/velox/exec/RowContainer.cpp
@@ -567,7 +567,8 @@ void RowContainer::store(
         decoded,
         rows,
         isKey,
-        offsets_[column]);
+        offsets_[column],
+        column);
   } else {
     const auto rowColumn = rowColumns_[column];
     VELOX_DYNAMIC_TYPE_DISPATCH_ALL(

diff --git a/velox/exec/RowContainer.h b/velox/exec/RowContainer.h
@@ -1056,6 +1056,7 @@ class RowContainer {
     for (int32_t i = 0; i < rows.size(); ++i) {
       storeWithNulls<Kind>(
           decoded, i, isKey, rows[i], offset, nullByte, nullMask, column);
+      updateColumnStats(decoded, i, rows[i], column);
     }
   }
 
@@ -1064,9 +1065,11 @@ class RowContainer {
       const DecodedVector& decoded,
       folly::Range<char**> rows,
       bool isKey,
-      int32_t offset) {
+      int32_t offset,
+      int32_t column) {
     for (int32_t i = 0; i < rows.size(); ++i) {
       storeNoNulls<Kind>(decoded, i, isKey, rows[i], offset);
+      updateColumnStats(decoded, i, rows[i], column);
     }
   }
 

diff --git a/velox/exec/Window.cpp b/velox/exec/Window.cpp
@@ -60,7 +60,8 @@ Window::Window(
         pool(),
         common::PrefixSortConfig{
             driverCtx->queryConfig().prefixSortNormalizedKeyMaxBytes(),
-            driverCtx->queryConfig().prefixSortMinRows()},
+            driverCtx->queryConfig().prefixSortMinRows(),
+            driverCtx->queryConfig().prefixSortMaxStringPrefixLength()},
         spillConfig,
         &nonReclaimableSection_,
         &spillStats_);

diff --git a/velox/exec/benchmarks/PrefixSortBenchmark.cpp b/velox/exec/benchmarks/PrefixSortBenchmark.cpp
@@ -125,15 +125,14 @@ class TestCase {
 
 // You could config threshold, e.i. 0, to test prefix-sort for small
 // dateset.
-static const common::PrefixSortConfig kDefaultSortConfig(1024, 100);
+static const common::PrefixSortConfig kDefaultSortConfig(1024, 100, 50);
 
 // For small dataset, in some test environments, if std-sort is defined in the
 // benchmark file, the test results may be strangely regressed. When the
 // threshold is particularly large, PrefixSort is actually std-sort, hence, we
 // can use this as std-sort benchmark base.
-static const common::PrefixSortConfig kStdSortConfig(
-    1024,
-    std::numeric_limits<int>::max());
+static const common::PrefixSortConfig
+    kStdSortConfig(1024, std::numeric_limits<int>::max(), 50);
 
 class PrefixSortBenchmark {
  public: