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

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

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);
   }

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:
 

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()};
   }
 };
 

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);

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;

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(

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);
     }
   }
 

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_);

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: