feat(prefix sort): Exclude null byte from sort prefix if column has n…

…o nulls (#11583)

Summary:
When inserting data into a RowContainer, we collect information about whether
columns contain null values. If a column does not contain nulls, we can omit
the null byte in the normalized key. This reduces memory usage and decreases
word-by-word comparisons.

For example, with a single sort key of type bigint, without this feature, the
normalized key size is 16 bytes (8 + 1 + alignment padding). With this feature,
the normalized key size is reduced to 8 bytes. However, this improvement only
makes a difference if alignment padding does not negate the optimization. For
instance, with a single sort key of type int, the normalized key size remains 8
bytes regardless of whether this feature is enabled.

Pull Request resolved: #11583

Reviewed By: tanjialiang

Differential Revision: D67587505

Pulled By: xiaoxmeng

fbshipit-source-id: 01d7d44d10079f5a4a389518a39eea2ae9c8fab0

velox/exec/PrefixSort.cpp

@@ -34,15 +34,18 @@ FOLLY_ALWAYS_INLINE void encodeRowColumn(
     char* row,
     char* prefixBuffer) {
   std::optional<T> value;
-  if (RowContainer::isNullAt(row, rowColumn.nullByte(), rowColumn.nullMask())) {
-    value = std::nullopt;
-  } else {
+  if (!prefixSortLayout.normalizedKeyHasNullByte[index] ||
+      !RowContainer::isNullAt(
+          row, rowColumn.nullByte(), rowColumn.nullMask())) {
     value = *(reinterpret_cast<T*>(row + rowColumn.offset()));
+  } else {
+    value = std::nullopt;
   }
   prefixSortLayout.encoders[index].encode(
       value,
       prefixBuffer + prefixSortLayout.prefixOffsets[index],
-      prefixSortLayout.encodeSizes[index]);
+      prefixSortLayout.encodeSizes[index],
+      prefixSortLayout.normalizedKeyHasNullByte[index]);
 }
 
 FOLLY_ALWAYS_INLINE void extractRowColumnToPrefix(
@@ -138,11 +141,14 @@ compareByWord(uint64_t* left, uint64_t* right, int32_t bytes) {
 // static.
 PrefixSortLayout PrefixSortLayout::generate(
     const std::vector<TypePtr>& types,
+    const std::vector<bool>& columnHasNulls,
     const std::vector<CompareFlags>& compareFlags,
     uint32_t maxNormalizedKeySize,
     uint32_t maxStringPrefixLength,
     const std::vector<std::optional<uint32_t>>& maxStringLengths) {
   const uint32_t numKeys = types.size();
+  std::vector<bool> normalizedKeyHasNullByte;
+  normalizedKeyHasNullByte.reserve(numKeys);
   std::vector<uint32_t> prefixOffsets;
   prefixOffsets.reserve(numKeys);
   std::vector<uint32_t> encodeSizes;
@@ -162,14 +168,16 @@ PrefixSortLayout PrefixSortLayout::generate(
         types[i]->kind(),
         maxStringLengths[i].has_value()
             ? std::min(maxStringLengths[i].value(), maxStringPrefixLength)
-            : maxStringPrefixLength);
+            : maxStringPrefixLength,
+        columnHasNulls[i]);
     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());
+    normalizedKeyHasNullByte.push_back(columnHasNulls[i]);
     normalizedKeySize += encodedSize.value();
     ++numNormalizedKeys;
     if ((types[i]->kind() == TypeKind::VARCHAR ||
@@ -194,6 +202,7 @@ PrefixSortLayout PrefixSortLayout::generate(
       numNormalizedKeys < numKeys,
       /*nonPrefixSortStartIndex=*/
       lastPrefixKeyPartial ? numNormalizedKeys - 1 : numNormalizedKeys,
+      std::move(normalizedKeyHasNullByte),
       std::move(prefixOffsets),
       std::move(encodeSizes),
       std::move(encoders),
@@ -210,7 +219,7 @@ void PrefixSortLayout::optimizeSortKeysOrder(
     // 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(), UINT_MAX - 1);
+        rowType->childAt(projection.inputChannel)->kind(), UINT_MAX - 1, true);
   }
 
   std::sort(

velox/exec/PrefixSort.h

@@ -60,6 +60,9 @@ struct PrefixSortLayout {
   /// numNormalizedKeys - 1. Otherwise, start from numNormalizedKeys.
   const uint32_t nonPrefixSortStartIndex;
 
+  /// A vector indicating whether each normalized key contains a null byte.
+  const std::vector<bool> normalizedKeyHasNullByte;
+
   /// Offsets of normalized keys, used to find write locations when
   /// extracting columns
   const std::vector<uint32_t> prefixOffsets;
@@ -76,6 +79,7 @@ struct PrefixSortLayout {
 
   static PrefixSortLayout generate(
       const std::vector<TypePtr>& types,
+      const std::vector<bool>& columnHasNulls,
       const std::vector<CompareFlags>& compareFlags,
       uint32_t maxNormalizedKeySize,
       uint32_t maxStringPrefixLength,
@@ -176,6 +180,8 @@ class PrefixSort {
     VELOX_CHECK_EQ(keyTypes.size(), compareFlags.size());
     std::vector<std::optional<uint32_t>> maxStringLengths;
     maxStringLengths.reserve(keyTypes.size());
+    std::vector<bool> columnHasNulls;
+    columnHasNulls.reserve(keyTypes.size());
     for (int i = 0; i < keyTypes.size(); ++i) {
       std::optional<uint32_t> maxStringLength = std::nullopt;
       if (keyTypes[i]->kind() == TypeKind::VARBINARY ||
@@ -186,9 +192,11 @@ class PrefixSort {
         }
       }
       maxStringLengths.emplace_back(maxStringLength);
+      columnHasNulls.emplace_back(rowContainer->columnHasNulls(i));
     }
     return PrefixSortLayout::generate(
         keyTypes,
+        columnHasNulls,
         compareFlags,
         config.maxNormalizedKeyBytes,
         config.maxStringPrefixLength,

velox/exec/benchmarks/PrefixSortBenchmark.cpp

@@ -20,6 +20,8 @@
 #include "velox/exec/PrefixSort.h"
 #include "velox/vector/fuzzer/VectorFuzzer.h"
 
+DEFINE_double(data_null_ratio, 0.7, "Data null ratio");
+
 using namespace facebook::velox;
 using namespace facebook::velox::exec;
 
@@ -92,7 +94,7 @@ class TestCase {
   RowVectorPtr fuzzRows(size_t numRows, int numKeys) {
     VectorFuzzer fuzzer({.vectorSize = numRows}, pool_);
     VectorFuzzer fuzzerWithNulls(
-        {.vectorSize = numRows, .nullRatio = 0.7}, pool_);
+        {.vectorSize = numRows, .nullRatio = FLAGS_data_null_ratio}, pool_);
     std::vector<VectorPtr> children;
 
     // Fuzz keys: for front keys (column 0 to numKeys -2) use high

velox/exec/prefixsort/PrefixSortEncoder.h

@@ -34,15 +34,21 @@ class PrefixSortEncoder {
   /// int32_t, uint16_t, int16_t, float, double, Timestamp).
   /// 1. The first byte of the encoded result is null byte. The value is 0 if
   ///    (nulls first and value is null) or (nulls last and value is not null).
-  ///    Otherwise, the value is 1.
+  ///    Otherwise, the value is 1. If this column has no null values, we can
+  ///    skip the null byte.
   /// 2. The remaining bytes are the encoding result of value:
   ///    -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.
   template <typename T>
-  FOLLY_ALWAYS_INLINE void
-  encode(const std::optional<T>& value, char* dest, uint32_t encodeSize) const {
-    if (value.has_value()) {
+  FOLLY_ALWAYS_INLINE void encode(
+      std::optional<T> value,
+      char* dest,
+      uint32_t encodeSize,
+      bool includeNullByte) const {
+    if (!includeNullByte) {
+      encodeNoNulls(value.value(), dest, encodeSize);
+    } else if (value.has_value()) {
       dest[0] = nullsFirst_ ? 1 : 0;
       encodeNoNulls(value.value(), dest + 1, encodeSize - 1);
     } else {
@@ -69,34 +75,27 @@ class PrefixSortEncoder {
   ///         For not supported types, returns 'std::nullopt'.
   FOLLY_ALWAYS_INLINE static std::optional<uint32_t> encodedSize(
       TypeKind typeKind,
-      uint32_t maxStringPrefixLength) {
-    // NOTE: one byte is reserved for nullable comparison.
+      uint32_t maxStringPrefixLength,
+      bool columnHasNulls) {
+    // NOTE: if columnHasNulls is true, one byte is reserved for nullable
+    // comparison.
+    const uint32_t nullByteSize = columnHasNulls ? 1 : 0;
     switch ((typeKind)) {
-      case ::facebook::velox::TypeKind::SMALLINT: {
-        return 3;
-      }
-      case ::facebook::velox::TypeKind::INTEGER: {
-        return 5;
-      }
-      case ::facebook::velox::TypeKind::BIGINT: {
-        return 9;
-      }
-      case ::facebook::velox::TypeKind::REAL: {
-        return 5;
-      }
-      case ::facebook::velox::TypeKind::DOUBLE: {
-        return 9;
-      }
-      case ::facebook::velox::TypeKind::TIMESTAMP: {
-        return 17;
-      }
-      case ::facebook::velox::TypeKind::HUGEINT: {
-        return 17;
-      }
-      case ::facebook::velox::TypeKind::VARBINARY:
+#define SCALAR_CASE(kind) \
+  case TypeKind::kind:    \
+    return nullByteSize + sizeof(TypeTraits<TypeKind::kind>::NativeType);
+      SCALAR_CASE(SMALLINT)
+      SCALAR_CASE(INTEGER)
+      SCALAR_CASE(BIGINT)
+      SCALAR_CASE(HUGEINT)
+      SCALAR_CASE(REAL)
+      SCALAR_CASE(DOUBLE)
+      SCALAR_CASE(TIMESTAMP)
+#undef SCALAR_CASE
+      case TypeKind::VARBINARY:
         [[fallthrough]];
-      case ::facebook::velox::TypeKind::VARCHAR: {
-        return 1 + maxStringPrefixLength;
+      case TypeKind::VARCHAR: {
+        return nullByteSize + maxStringPrefixLength;
       }
       default:
         return std::nullopt;