Support string type key in prefix sort

velox/common/base/PrefixSortConfig.h

@@ -24,14 +24,23 @@ namespace facebook::velox::common {
 struct PrefixSortConfig {
   PrefixSortConfig() = default;
 
-  PrefixSortConfig(int64_t _maxNormalizedKeySize, int32_t _threshold)
-      : maxNormalizedKeySize(_maxNormalizedKeySize), threshold(_threshold) {}
+  PrefixSortConfig(
+      int64_t _maxNormalizedKeySize,
+      int32_t _threshold,
+      int32_t _stringPrefixLength)
+      : maxNormalizedKeySize(_maxNormalizedKeySize),
+        threshold(_threshold),
+        stringPrefixLength(_stringPrefixLength) {}
 
   /// Max number of bytes can store normalized keys in prefix-sort buffer per
   /// entry. Same with QueryConfig kPrefixSortNormalizedKeyMaxBytes.
   int64_t maxNormalizedKeySize{128};
 
   /// PrefixSort will have performance regression when the dateset is too small.
   int32_t threshold{130};
+
+  /// Length of the prefix to be stored in prefix-sort buffer for a string
+  /// column.
+  int32_t stringPrefixLength{12};
 };
 } // 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";
 
+  /// Length of the prefix to be stored in prefix-sort buffer for a string
+  /// key.
+  static constexpr const char* kPrefixSortStringPrefixLength =
+      "prefixsort_string_prefix_length";
+
   /// Enable query tracing flag.
   static constexpr const char* kQueryTraceEnabled = "query_trace_enabled";
 
@@ -796,6 +801,10 @@ class QueryConfig {
     return get<int32_t>(kPrefixSortMinRows, 130);
   }
 
+  int32_t prefixSortStringPrefixLength() const {
+    return get<int32_t>(kPrefixSortStringPrefixLength, 12);
+  }
+
   template <typename T>
   T get(const std::string& key, const T& defaultValue) const {
     return config_->get<T>(key, defaultValue);

velox/docs/configs.rst

@@ -144,6 +144,10 @@ Generic Configuration
      - integer
      - 130
      - Minimum number of rows to use prefix-sort. The default value has been derived using micro-benchmarking.
+   * - prefixsort_string_prefix_length
+     - integer
+     - 12
+     - 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

@@ -303,7 +303,8 @@ struct DriverCtx {
   common::PrefixSortConfig prefixSortConfig() const {
     return common::PrefixSortConfig{
         queryConfig().prefixSortNormalizedKeyMaxBytes(),
-        queryConfig().prefixSortMinRows()};
+        queryConfig().prefixSortMinRows(),
+        queryConfig().prefixSortStringPrefixLength()};
   }
 };
 

velox/exec/PrefixSort.cpp

@@ -39,8 +39,15 @@ FOLLY_ALWAYS_INLINE void encodeRowColumn(
   } else {
     value = *(reinterpret_cast<T*>(row + rowColumn.offset()));
   }
-  prefixSortLayout.encoders[index].encode(
-      value, prefixBuffer + prefixSortLayout.prefixOffsets[index]);
+  if constexpr (std::is_same_v<T, StringView>) {
+    prefixSortLayout.encoders[index].encode(
+        value,
+        prefixBuffer + prefixSortLayout.prefixOffsets[index],
+        prefixSortLayout.encodeSizes[index]);
+  } else {
+    prefixSortLayout.encoders[index].encode(
+        value, prefixBuffer + prefixSortLayout.prefixOffsets[index]);
+  }
 }
 
 FOLLY_ALWAYS_INLINE void extractRowColumnToPrefix(
@@ -86,6 +93,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: {}",
@@ -129,28 +143,43 @@ compareByWord(uint64_t* left, uint64_t* right, int32_t bytes) {
 PrefixSortLayout PrefixSortLayout::makeSortLayout(
     const std::vector<TypePtr>& types,
     const std::vector<CompareFlags>& compareFlags,
-    uint32_t maxNormalizedKeySize) {
+    uint32_t maxNormalizedKeySize,
+    int32_t stringPrefixLength) {
   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 lastKeyInPrefixIsPartial{false};
   for (auto i = 0; i < numKeys; ++i) {
     const std::optional<uint32_t> encodedSize =
-        PrefixSortEncoder::encodedSize(types[i]->kind());
+        PrefixSortEncoder::encodedSize(types[i]->kind(), stringPrefixLength);
     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;
+    // Since we can't be certain that the maximum length of the string keys is
+    // <= 'encodedSize', we can only assume that partial data will be stored in
+    // the prefix and stop the loop.
+    if (types[i]->kind() == TypeKind::VARCHAR ||
+        types[i]->kind() == TypeKind::VARBINARY) {
+      lastKeyInPrefixIsPartial = true;
+      break;
+    }
   }
 
   const auto numPaddingBytes = alignmentPadding(normalizedKeySize, kAlignment);
@@ -163,8 +192,9 @@ PrefixSortLayout PrefixSortLayout::makeSortLayout(
       numKeys,
       compareFlags,
       numNormalizedKeys != 0,
-      numNormalizedKeys < numKeys,
+      numNormalizedKeys < numKeys || lastKeyInPrefixIsPartial,
       std::move(prefixOffsets),
+      std::move(encodeSizes),
       std::move(encoders),
       numPaddingBytes};
 }
@@ -241,7 +271,10 @@ uint32_t PrefixSort::maxRequiredBytes(
   }
   VELOX_CHECK_EQ(rowContainer->keyTypes().size(), compareFlags.size());
   const auto sortLayout = PrefixSortLayout::makeSortLayout(
-      rowContainer->keyTypes(), compareFlags, config.maxNormalizedKeySize);
+      rowContainer->keyTypes(),
+      compareFlags,
+      config.maxNormalizedKeySize,
+      config.stringPrefixLength);
   if (!sortLayout.hasNormalizedKeys) {
     return 0;
   }
@@ -303,7 +336,7 @@ void PrefixSort::sortInternal(
     PrefixSortRunner sortRunner(entrySize, swapBuffer->asMutable<char>());
     auto* prefixBufferStart = prefixBuffer;
     auto* prefixBufferEnd = prefixBuffer + numRows * entrySize;
-    if (sortLayout_.hasNonNormalizedKey) {
+    if (sortLayout_.hasNonFullyNormalizedKey) {
       sortRunner.quickSort(
           prefixBufferStart, prefixBufferEnd, [&](char* lhs, char* rhs) {
             return comparePartNormalizedKeys(lhs, rhs);

velox/exec/PrefixSort.h

@@ -49,13 +49,16 @@ struct PrefixSortLayout {
   /// It equals to 'numNormalizedKeys != 0', a little faster.
   const bool hasNormalizedKeys;
 
-  /// Whether the sort keys contains non-normalized key.
-  const bool hasNonNormalizedKey;
+  /// Whether the sort keys contains non-normalized or partial normalized key.
+  const bool hasNonFullyNormalizedKey;
 
   /// 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;
 
@@ -66,7 +69,8 @@ struct PrefixSortLayout {
   static PrefixSortLayout makeSortLayout(
       const std::vector<TypePtr>& types,
       const std::vector<CompareFlags>& compareFlags,
-      uint32_t maxNormalizedKeySize);
+      uint32_t maxNormalizedKeySize,
+      int32_t stringPrefixLength);
 };
 
 class PrefixSort {
@@ -112,7 +116,10 @@ class PrefixSort {
 
     VELOX_CHECK_EQ(rowContainer->keyTypes().size(), compareFlags.size());
     const auto sortLayout = PrefixSortLayout::makeSortLayout(
-        rowContainer->keyTypes(), compareFlags, config.maxNormalizedKeySize);
+        rowContainer->keyTypes(),
+        compareFlags,
+        config.maxNormalizedKeySize,
+        config.stringPrefixLength);
     // All keys can not normalize, skip the binary string compare opt.
     // Putting this outside sort-internal helps with stdSort.
     if (!sortLayout.hasNormalizedKeys) {

velox/exec/Window.cpp

@@ -60,7 +60,8 @@ Window::Window(
         pool(),
         common::PrefixSortConfig{
             driverCtx->queryConfig().prefixSortNormalizedKeyMaxBytes(),
-            driverCtx->queryConfig().prefixSortMinRows()},
+            driverCtx->queryConfig().prefixSortMinRows(),
+            driverCtx->queryConfig().prefixSortStringPrefixLength()},
         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, 12);
 
 // 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(), 12);
 
 class PrefixSortBenchmark {
  public:

velox/exec/prefixsort/PrefixSortEncoder.h

@@ -18,6 +18,7 @@
 #include <cstdint>
 
 #include "velox/common/base/SimdUtil.h"
+#include "velox/common/memory/HashStringAllocator.h"
 #include "velox/type/Timestamp.h"
 #include "velox/type/Type.h"
 
@@ -38,8 +39,6 @@ class PrefixSortEncoder {
   ///    -If value is null, we set the remaining sizeof(T) bytes to '0', they
   ///     do not affect the comparison results at all.
   ///    -If value is not null, the result is set by calling encodeNoNulls.
-  ///
-  /// TODO: add support for strings.
   template <typename T>
   FOLLY_ALWAYS_INLINE void encode(std::optional<T> value, char* dest) const {
     if (value.has_value()) {
@@ -51,6 +50,49 @@ class PrefixSortEncoder {
     }
   }
 
+  /// Encodes String types.
+  /// The string prefix is formatted as 'null byte + string content + padding
+  /// zeros'. If `!ascending_`, the bits for both the content and padding zeros
+  /// need to be inverted.
+  FOLLY_ALWAYS_INLINE void encode(
+      std::optional<StringView> value,
+      char* dest,
+      int32_t encodeSize) const {
+    auto* destDataPtr = dest + 1;
+    const auto stringPrefixSize = encodeSize - 1;
+    if (value.has_value()) {
+      dest[0] = nullsFirst_ ? 1 : 0;
+      auto data = value.value();
+      const int32_t copySize = std::min<int32_t>(data.size(), stringPrefixSize);
+      if (data.isInline() ||
+          reinterpret_cast<const HashStringAllocator::Header*>(data.data())[-1]
+                  .size() >= data.size()) {
+        // The string is inline or all in one piece out of line.
+        std::memcpy(destDataPtr, data.data(), copySize);
+      } else {
+        // 'data' is stored in non-contiguous allocation pieces in the row
+        // container, we only read prefix size data out.
+        auto stream = HashStringAllocator::prepareRead(
+            HashStringAllocator::headerOf(data.data()));
+        stream->readBytes(destDataPtr, copySize);
+      }
+
+      if (data.size() < stringPrefixSize) {
+        std::memset(
+            destDataPtr + data.size(), 0, stringPrefixSize - data.size());
+      }
+
+      if (!ascending_) {
+        for (auto i = 1; i < encodeSize; ++i) {
+          dest[i] = ~dest[i];
+        }
+      }
+    } else {
+      dest[0] = nullsFirst_ ? 0 : 1;
+      std::memset(destDataPtr, 0, stringPrefixSize);
+    }
+  }
+
   /// @tparam T Type of value. Supported type are: uint64_t, int64_t, uint32_t,
   /// int32_t, int16_t, uint16_t, float, double, Timestamp.
   template <typename T>
@@ -67,7 +109,8 @@ class PrefixSortEncoder {
   /// @return For supported types, returns the encoded size, assume nullable.
   ///         For not supported types, returns 'std::nullopt'.
   FOLLY_ALWAYS_INLINE static std::optional<uint32_t> encodedSize(
-      TypeKind typeKind) {
+      TypeKind typeKind,
+      int32_t stringPrefixLength) {
     // NOTE: one byte is reserved for nullable comparison.
     switch ((typeKind)) {
       case ::facebook::velox::TypeKind::SMALLINT: {
@@ -91,6 +134,11 @@ class PrefixSortEncoder {
       case ::facebook::velox::TypeKind::HUGEINT: {
         return 17;
       }
+      case ::facebook::velox::TypeKind::VARBINARY:
+        [[fallthrough]];
+      case ::facebook::velox::TypeKind::VARCHAR: {
+        return 1 + stringPrefixLength;
+      }
       default:
         return std::nullopt;
     }
@@ -260,4 +308,11 @@ FOLLY_ALWAYS_INLINE void PrefixSortEncoder::encodeNoNulls(
   encodeNoNulls(value.getNanos(), dest + 8);
 }
 
+template <>
+FOLLY_ALWAYS_INLINE void PrefixSortEncoder::encodeNoNulls(
+    StringView value,
+    char* dest) const {
+  VELOX_UNREACHABLE();
+}
+
 } // namespace facebook::velox::exec::prefixsort