diff --git a/velox/common/base/PrefixSortConfig.h b/velox/common/base/PrefixSortConfig.h index 46ac49391554..d53f44d6486f 100644 --- 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 index 7bceb06d117b..ef6982217270 100644 --- 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(kPrefixSortMinRows, 128); } + uint32_t prefixSortMaxStringPrefixLength() const { + return get(kPrefixSortMaxStringPrefixLength, 16); + } + double scaleWriterRebalanceMaxMemoryUsageRatio() const { return get(kScaleWriterRebalanceMaxMemoryUsageRatio, 0.7); } diff --git a/velox/docs/configs.rst b/velox/docs/configs.rst index af7e26a60abb..e7a669c89ae5 100644 --- 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 index e7ae1782a4a0..7d46017187e4 100644 --- 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 index b4dfb1f6e0d1..e7ef749d2edd 100644 --- a/velox/exec/PrefixSort.cpp +++ b/velox/exec/PrefixSort.cpp @@ -40,7 +40,9 @@ FOLLY_ALWAYS_INLINE void encodeRowColumn( value = *(reinterpret_cast(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( + 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& types, const std::vector& compareFlags, - uint32_t maxNormalizedKeySize) { + uint32_t maxNormalizedKeySize, + uint32_t maxStringPrefixLength, + const std::vector>& maxStringLengths) { const uint32_t numKeys = types.size(); std::vector prefixOffsets; prefixOffsets.reserve(numKeys); + std::vector encodeSizes; + encodeSizes.reserve(numKeys); std::vector 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 encodedSize = - PrefixSortEncoder::encodedSize(types[i]->kind()); + const std::optional 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> 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 index 7ac34b430551..292bdbe524bc 100644 --- a/velox/exec/PrefixSort.h +++ b/velox/exec/PrefixSort.h @@ -15,6 +15,8 @@ */ #pragma once +#include + #include "velox/common/base/PrefixSortConfig.h" #include "velox/exec/Operator.h" #include "velox/exec/RowContainer.h" @@ -53,10 +55,18 @@ 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 prefixOffsets; + /// Sizes of normalized keys. + const std::vector encodeSizes; + /// The encoders for normalized keys. const std::vector encoders; @@ -64,10 +74,12 @@ struct PrefixSortLayout { /// for fast long compare. const int32_t numPaddingBytes; - static PrefixSortLayout makeSortLayout( + static PrefixSortLayout generate( const std::vector& types, const std::vector& compareFlags, - uint32_t maxNormalizedKeySize); + uint32_t maxNormalizedKeySize, + uint32_t maxStringPrefixLength, + const std::vector>& 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); + FOLLY_ALWAYS_INLINE static PrefixSortLayout generateSortLayout( + const RowContainer* rowContainer, + const std::vector& compareFlags, + const velox::common::PrefixSortConfig& config) { + const auto keyTypes = rowContainer->keyTypes(); + VELOX_CHECK_EQ(keyTypes.size(), compareFlags.size()); + std::vector> maxStringLengths; + maxStringLengths.reserve(keyTypes.size()); + for (int i = 0; i < keyTypes.size(); ++i) { + std::optional 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 index d8b622c3a728..74d49266b0fd 100644 --- 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 index 3782dff69816..72cd08e276dd 100644 --- 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( 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 rows, bool isKey, - int32_t offset) { + int32_t offset, + int32_t column) { for (int32_t i = 0; i < rows.size(); ++i) { storeNoNulls(decoded, i, isKey, rows[i], offset); + updateColumnStats(decoded, i, rows[i], column); } } diff --git a/velox/exec/Window.cpp b/velox/exec/Window.cpp index c8694cff9574..1f1a91331a95 100644 --- 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 index 692d4f72926b..12e575b68e80 100644 --- 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::max()); +static const common::PrefixSortConfig + kStdSortConfig(1024, std::numeric_limits::max(), 50); class PrefixSortBenchmark { public: diff --git a/velox/exec/prefixsort/PrefixSortEncoder.h b/velox/exec/prefixsort/PrefixSortEncoder.h index c4d356ec53d4..e396a980171f 100644 --- a/velox/exec/prefixsort/PrefixSortEncoder.h +++ b/velox/exec/prefixsort/PrefixSortEncoder.h @@ -18,6 +18,7 @@ #include #include "velox/common/base/SimdUtil.h" +#include "velox/common/memory/HashStringAllocator.h" #include "velox/type/Timestamp.h" #include "velox/type/Type.h" @@ -30,7 +31,7 @@ class PrefixSortEncoder { : ascending_(ascending), nullsFirst_(nullsFirst){}; /// Encode native primitive types(such as uint64_t, int64_t, uint32_t, - /// int32_t, float, double, Timestamp). + /// 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. @@ -38,23 +39,23 @@ 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 - FOLLY_ALWAYS_INLINE void encode(std::optional value, char* dest) const { + FOLLY_ALWAYS_INLINE void + encode(const std::optional& value, char* dest, uint32_t encodeSize) const { if (value.has_value()) { dest[0] = nullsFirst_ ? 1 : 0; - encodeNoNulls(value.value(), dest + 1); + encodeNoNulls(value.value(), dest + 1, encodeSize - 1); } else { dest[0] = nullsFirst_ ? 0 : 1; - simd::memset(dest + 1, 0, sizeof(T)); + simd::memset(dest + 1, 0, encodeSize - 1); } } /// @tparam T Type of value. Supported type are: uint64_t, int64_t, uint32_t, /// int32_t, int16_t, uint16_t, float, double, Timestamp. template - FOLLY_ALWAYS_INLINE void encodeNoNulls(T value, char* dest) const; + FOLLY_ALWAYS_INLINE void + encodeNoNulls(T value, char* dest, uint32_t encodeSize) const; bool isAscending() const { return ascending_; @@ -67,7 +68,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 encodedSize( - TypeKind typeKind) { + TypeKind typeKind, + uint32_t maxStringPrefixLength) { // NOTE: one byte is reserved for nullable comparison. switch ((typeKind)) { case ::facebook::velox::TypeKind::SMALLINT: { @@ -91,6 +93,11 @@ class PrefixSortEncoder { case ::facebook::velox::TypeKind::HUGEINT: { return 17; } + case ::facebook::velox::TypeKind::VARBINARY: + [[fallthrough]]; + case ::facebook::velox::TypeKind::VARCHAR: { + return 1 + maxStringPrefixLength; + } default: return std::nullopt; } @@ -112,7 +119,8 @@ class PrefixSortEncoder { template <> FOLLY_ALWAYS_INLINE void PrefixSortEncoder::encodeNoNulls( uint32_t value, - char* dest) const { + char* dest, + uint32_t /*encodeSize*/) const { auto& v = *reinterpret_cast(dest); v = __builtin_bswap32(value); if (!ascending_) { @@ -129,15 +137,17 @@ FOLLY_ALWAYS_INLINE void PrefixSortEncoder::encodeNoNulls( template <> FOLLY_ALWAYS_INLINE void PrefixSortEncoder::encodeNoNulls( int32_t value, - char* dest) const { - encodeNoNulls((uint32_t)(value ^ (1u << 31)), dest); + char* dest, + uint32_t encodeSize) const { + encodeNoNulls((uint32_t)(value ^ (1u << 31)), dest, encodeSize); } /// Logic is as same as int32_t. template <> FOLLY_ALWAYS_INLINE void PrefixSortEncoder::encodeNoNulls( uint64_t value, - char* dest) const { + char* dest, + uint32_t /*encodeSize*/) const { auto& v = *reinterpret_cast(dest); v = __builtin_bswap64(value); if (!ascending_) { @@ -148,15 +158,17 @@ FOLLY_ALWAYS_INLINE void PrefixSortEncoder::encodeNoNulls( template <> FOLLY_ALWAYS_INLINE void PrefixSortEncoder::encodeNoNulls( int64_t value, - char* dest) const { - encodeNoNulls((uint64_t)(value ^ (1ull << 63)), dest); + char* dest, + uint32_t encodeSize) const { + encodeNoNulls((uint64_t)(value ^ (1ull << 63)), dest, encodeSize); } /// Logic is as same as int32_t. template <> FOLLY_ALWAYS_INLINE void PrefixSortEncoder::encodeNoNulls( uint16_t value, - char* dest) const { + char* dest, + uint32_t /*encodeSize*/) const { auto& v = *reinterpret_cast(dest); v = __builtin_bswap16(value); if (!ascending_) { @@ -167,16 +179,19 @@ FOLLY_ALWAYS_INLINE void PrefixSortEncoder::encodeNoNulls( template <> FOLLY_ALWAYS_INLINE void PrefixSortEncoder::encodeNoNulls( int16_t value, - char* dest) const { - encodeNoNulls(static_cast(value ^ (1u << 15)), dest); + char* dest, + uint32_t encodeSize) const { + encodeNoNulls(static_cast(value ^ (1u << 15)), dest, encodeSize); } template <> FOLLY_ALWAYS_INLINE void PrefixSortEncoder::encodeNoNulls( int128_t value, - char* dest) const { - encodeNoNulls(HugeInt::upper(value), dest); - encodeNoNulls(HugeInt::lower(value), dest + sizeof(int64_t)); + char* dest, + uint32_t encodeSize) const { + encodeNoNulls(HugeInt::upper(value), dest, encodeSize); + encodeNoNulls( + HugeInt::lower(value), dest + sizeof(int64_t), encodeSize); } namespace detail { @@ -239,15 +254,17 @@ static FOLLY_ALWAYS_INLINE uint32_t encodeFloat(float value) { template <> FOLLY_ALWAYS_INLINE void PrefixSortEncoder::encodeNoNulls( double value, - char* dest) const { - encodeNoNulls(detail::encodeDouble(value), dest); + char* dest, + uint32_t encodeSize) const { + encodeNoNulls(detail::encodeDouble(value), dest, encodeSize); } template <> FOLLY_ALWAYS_INLINE void PrefixSortEncoder::encodeNoNulls( float value, - char* dest) const { - encodeNoNulls(detail::encodeFloat(value), dest); + char* dest, + uint32_t encodeSize) const { + encodeNoNulls(detail::encodeFloat(value), dest, encodeSize); } /// When comparing Timestamp, first compare seconds and then compare nanos, so @@ -255,9 +272,43 @@ FOLLY_ALWAYS_INLINE void PrefixSortEncoder::encodeNoNulls( template <> FOLLY_ALWAYS_INLINE void PrefixSortEncoder::encodeNoNulls( Timestamp value, - char* dest) const { - encodeNoNulls(value.getSeconds(), dest); - encodeNoNulls(value.getNanos(), dest + 8); + char* dest, + uint32_t encodeSize) const { + encodeNoNulls(value.getSeconds(), dest, encodeSize); + encodeNoNulls(value.getNanos(), dest + 8, encodeSize); +} + +/// Encode String type. +/// 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. +template <> +FOLLY_ALWAYS_INLINE void PrefixSortEncoder::encodeNoNulls( + StringView value, + char* dest, + uint32_t encodeSize) const { + const uint32_t copySize = std::min(value.size(), encodeSize); + if (value.isInline() || + HashStringAllocator::headerOf(value.data())->size() >= value.size()) { + // The string is inline or all in one piece out of line. + std::memcpy(dest, value.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(value.data())); + stream->readBytes(dest, copySize); + } + + if (value.size() < encodeSize) { + std::memset(dest + value.size(), 0, encodeSize - value.size()); + } + + if (!ascending_) { + for (auto i = 0; i < encodeSize; ++i) { + dest[i] = ~dest[i]; + } + } } } // namespace facebook::velox::exec::prefixsort diff --git a/velox/exec/prefixsort/tests/PrefixEncoderTest.cpp b/velox/exec/prefixsort/tests/PrefixEncoderTest.cpp index dc21fabea7c1..072b4a822542 100644 --- a/velox/exec/prefixsort/tests/PrefixEncoderTest.cpp +++ b/velox/exec/prefixsort/tests/PrefixEncoderTest.cpp @@ -74,9 +74,9 @@ class PrefixEncoderTest : public testing::Test, template void testEncodeNoNull(T value, char* expectedAsc, char* expectedDesc) { char encoded[sizeof(T)]; - ascNullsFirstEncoder_.encodeNoNulls(value, (char*)encoded); + ascNullsFirstEncoder_.encodeNoNulls(value, (char*)encoded, sizeof(T) + 1); ASSERT_EQ(std::memcmp(encoded, expectedAsc, sizeof(T)), 0); - descNullsFirstEncoder_.encodeNoNulls(value, (char*)encoded); + descNullsFirstEncoder_.encodeNoNulls(value, (char*)encoded, sizeof(T) + 1); ASSERT_EQ(std::memcmp(encoded, expectedDesc, sizeof(T)), 0); } @@ -95,21 +95,21 @@ class PrefixEncoderTest : public testing::Test, return std::memcmp(left, right, sizeof(T) + 1); }; - ascNullsFirstEncoder_.encode(nullValue, encoded); + ascNullsFirstEncoder_.encode(nullValue, encoded, sizeof(T) + 1); ASSERT_EQ(compare(nullFirst, encoded), 0); - ascNullsLastEncoder_.encode(nullValue, encoded); + ascNullsLastEncoder_.encode(nullValue, encoded, sizeof(T) + 1); ASSERT_EQ(compare(nullLast, encoded), 0); - ascNullsFirstEncoder_.encode(value, encoded); + ascNullsFirstEncoder_.encode(value, encoded, sizeof(T) + 1); ASSERT_EQ(encoded[0], 1); ASSERT_EQ(std::memcmp(encoded + 1, expectedAsc, sizeof(T)), 0); - ascNullsLastEncoder_.encode(value, encoded); + ascNullsLastEncoder_.encode(value, encoded, sizeof(T) + 1); ASSERT_EQ(encoded[0], 0); ASSERT_EQ(std::memcmp(encoded + 1, expectedAsc, sizeof(T)), 0); - descNullsFirstEncoder_.encode(value, encoded); + descNullsFirstEncoder_.encode(value, encoded, sizeof(T) + 1); ASSERT_EQ(encoded[0], 1); ASSERT_EQ(std::memcmp(encoded + 1, expectedDesc, sizeof(T)), 0); - descNullsLastEncoder_.encode(value, encoded); + descNullsLastEncoder_.encode(value, encoded, sizeof(T) + 1); ASSERT_EQ(encoded[0], 0); ASSERT_EQ(std::memcmp(encoded + 1, expectedDesc, sizeof(T)), 0); } @@ -122,21 +122,22 @@ class PrefixEncoderTest : public testing::Test, template void testNullCompare() { + constexpr uint32_t kEncodeSize = sizeof(T) + 1; std::optional nullValue = std::nullopt; std::optional max = std::numeric_limits::max(); std::optional min = std::numeric_limits::min(); - char encodedNull[sizeof(T) + 1]; - char encodedMax[sizeof(T) + 1]; - char encodedMin[sizeof(T) + 1]; + char encodedNull[kEncodeSize]; + char encodedMax[kEncodeSize]; + char encodedMin[kEncodeSize]; auto encode = [&](auto& encoder) { - encoder.encode(nullValue, encodedNull); - encoder.encode(min, encodedMin); - encoder.encode(max, encodedMax); + encoder.encode(nullValue, encodedNull, kEncodeSize); + encoder.encode(min, encodedMin, kEncodeSize); + encoder.encode(max, encodedMax, kEncodeSize); }; auto compare = [](char* left, char* right) { - return std::memcmp(left, right, sizeof(T) + 1); + return std::memcmp(left, right, kEncodeSize); }; // Nulls first: NULL < non-NULL. @@ -154,34 +155,35 @@ class PrefixEncoderTest : public testing::Test, // For float / double`s NaN. if (TypeLimits::isFloat) { std::optional nan = TypeLimits::nan(); - char encodedNaN[sizeof(T) + 1]; + char encodedNaN[kEncodeSize]; - ascNullsFirstEncoder_.encode(nan, encodedNaN); - ascNullsFirstEncoder_.encode(max, encodedMax); + ascNullsFirstEncoder_.encode(nan, encodedNaN, kEncodeSize); + ascNullsFirstEncoder_.encode(max, encodedMax, kEncodeSize); ASSERT_GT(compare(encodedNaN, encodedMax), 0); - ascNullsFirstEncoder_.encode(nan, encodedNaN); - ascNullsFirstEncoder_.encode(nullValue, encodedNull); + ascNullsFirstEncoder_.encode(nan, encodedNaN, kEncodeSize); + ascNullsFirstEncoder_.encode(nullValue, encodedNull, kEncodeSize); ASSERT_LT(compare(encodedNull, encodedNaN), 0); } } template void testValidValueCompare() { + constexpr uint32_t kEncodeSize = sizeof(T) + 1; std::optional max = std::numeric_limits::max(); std::optional min = TypeLimits::min(); std::optional mid = TypeLimits::mid(); - char encodedMax[sizeof(T) + 1]; - char encodedMin[sizeof(T) + 1]; - char encodedMid[sizeof(T) + 1]; + char encodedMax[kEncodeSize]; + char encodedMin[kEncodeSize]; + char encodedMid[kEncodeSize]; auto encode = [&](auto& encoder) { - encoder.encode(mid, encodedMid); - encoder.encode(min, encodedMin); - encoder.encode(max, encodedMax); + encoder.encode(mid, encodedMid, kEncodeSize); + encoder.encode(min, encodedMin, kEncodeSize); + encoder.encode(max, encodedMax, kEncodeSize); }; auto compare = [](char* left, char* right) { - return std::memcmp(left, right, sizeof(T) + 1); + return std::memcmp(left, right, kEncodeSize); }; encode(ascNullsFirstEncoder_); @@ -225,7 +227,9 @@ class PrefixEncoderTest : public testing::Test, auto test = [&](const PrefixSortEncoder& encoder) { TypePtr type = TypeTraits::ImplType::create(); - VectorFuzzer fuzzer({.vectorSize = vectorSize, .nullRatio = 0.1}, pool()); + VectorFuzzer fuzzer( + {.vectorSize = vectorSize, .nullRatio = 0.1, .stringLength = 16}, + pool()); CompareFlags compareFlag = { encoder.isNullsFirst(), @@ -250,8 +254,14 @@ class PrefixEncoderTest : public testing::Test, const auto rightValue = rightVector->isNullAt(i) ? std::nullopt : std::optional(rightVector->valueAt(i)); - encoder.encode(leftValue, leftEncoded); - encoder.encode(rightValue, rightEncoded); + if constexpr ( + Kind == TypeKind::VARCHAR || Kind == TypeKind::VARBINARY) { + encoder.encode(leftValue, leftEncoded, 17); + encoder.encode(rightValue, rightEncoded, 17); + } else { + encoder.encode(leftValue, leftEncoded, sizeof(ValueDataType) + 1); + encoder.encode(rightValue, rightEncoded, sizeof(ValueDataType) + 1); + } const auto result = compare(leftEncoded, rightEncoded); const auto expected = @@ -264,7 +274,23 @@ class PrefixEncoderTest : public testing::Test, test(ascNullsLastEncoder_); test(descNullsFirstEncoder_); test(descNullsLastEncoder_); - }; + } + + const PrefixSortEncoder& ascNullsFirstEncoder() const { + return ascNullsFirstEncoder_; + } + + const PrefixSortEncoder ascNullsLastEncoder() const { + return ascNullsLastEncoder_; + } + + const PrefixSortEncoder descNullsFirstEncoder() const { + return descNullsFirstEncoder_; + } + + const PrefixSortEncoder descNullsLastEncoder() const { + return descNullsLastEncoder_; + } protected: static void SetUpTestCase() { @@ -347,6 +373,46 @@ TEST_F(PrefixEncoderTest, encode) { } } +TEST_F(PrefixEncoderTest, encodeString) { + constexpr uint32_t kEncodeSize = 13; + StringView testValue = StringView("aaaaaabbbbbb"); + char expectedAsc[kEncodeSize] = "aaaaaabbbbbb"; + char expectedDesc[kEncodeSize]; + for (int i = 0; i < kEncodeSize - 1; ++i) { + expectedDesc[i] = ~expectedAsc[i]; + } + std::optional nullValue = std::nullopt; + std::optional value = testValue; + char encoded[kEncodeSize + 1]; + char nullFirst[kEncodeSize + 1]; + char nullLast[kEncodeSize + 1]; + memset(nullFirst, 0, kEncodeSize); + memset(nullLast, 1, 1); + memset(nullLast + 1, 0, kEncodeSize - 1); + + auto compare = [&](char* left, char* right) { + return std::memcmp(left, right, kEncodeSize); + }; + + ascNullsFirstEncoder().encode(nullValue, encoded, kEncodeSize); + ASSERT_EQ(compare(nullFirst, encoded), 0); + ascNullsLastEncoder().encode(nullValue, encoded, kEncodeSize); + ASSERT_EQ(compare(nullLast, encoded), 0); + + ascNullsFirstEncoder().encode(value, encoded, kEncodeSize); + ASSERT_EQ(encoded[0], 1); + ASSERT_EQ(std::memcmp(encoded + 1, expectedAsc, kEncodeSize - 1), 0); + ascNullsLastEncoder().encode(value, encoded, kEncodeSize); + ASSERT_EQ(encoded[0], 0); + ASSERT_EQ(std::memcmp(encoded + 1, expectedAsc, kEncodeSize - 1), 0); + descNullsFirstEncoder().encode(value, encoded, kEncodeSize); + ASSERT_EQ(encoded[0], 1); + ASSERT_EQ(std::memcmp(encoded + 1, expectedDesc, kEncodeSize - 1), 0); + descNullsLastEncoder().encode(value, encoded, kEncodeSize); + ASSERT_EQ(encoded[0], 0); + ASSERT_EQ(std::memcmp(encoded + 1, expectedDesc, kEncodeSize - 1), 0); +} + TEST_F(PrefixEncoderTest, compare) { testCompare(); testCompare(); @@ -388,4 +454,12 @@ TEST_F(PrefixEncoderTest, fuzzyTimestamp) { testFuzz(); } +TEST_F(PrefixEncoderTest, fuzzyStringView) { + testFuzz(); +} + +TEST_F(PrefixEncoderTest, fuzzyBinary) { + testFuzz(); +} + } // namespace facebook::velox::exec::prefixsort::test diff --git a/velox/exec/prefixsort/tests/utils/EncoderTestUtils.cpp b/velox/exec/prefixsort/tests/utils/EncoderTestUtils.cpp index d1773def0a3f..75aba3163ff6 100644 --- a/velox/exec/prefixsort/tests/utils/EncoderTestUtils.cpp +++ b/velox/exec/prefixsort/tests/utils/EncoderTestUtils.cpp @@ -32,7 +32,8 @@ void decodeNoNulls(char* encoded, int64_t& value) { void encodeInPlace(std::vector& data) { const static auto encoder = PrefixSortEncoder(true, true); for (auto i = 0; i < data.size(); i++) { - encoder.encodeNoNulls(data[i], (char*)data.data() + i * sizeof(int64_t)); + encoder.encodeNoNulls( + data[i], (char*)data.data() + i * sizeof(int64_t), sizeof(int64_t) + 1); } } diff --git a/velox/exec/tests/AggregationTest.cpp b/velox/exec/tests/AggregationTest.cpp index d4feb22475c5..19827536cc28 100644 --- a/velox/exec/tests/AggregationTest.cpp +++ b/velox/exec/tests/AggregationTest.cpp @@ -2036,92 +2036,108 @@ TEST_F(AggregationTest, spillPrefixSortOptimization) { {true, 0, 0, 0}, {false, 0, 0, 0}, {true, - prefixsort::PrefixSortEncoder::encodedSize(TypeKind::SMALLINT).value() - + prefixsort::PrefixSortEncoder::encodedSize(TypeKind::SMALLINT, 12) + .value() - 1, 0, 0}, {false, - prefixsort::PrefixSortEncoder::encodedSize(TypeKind::SMALLINT).value() - + prefixsort::PrefixSortEncoder::encodedSize(TypeKind::SMALLINT, 12) + .value() - 1, 0, 0}, {true, - prefixsort::PrefixSortEncoder::encodedSize(TypeKind::SMALLINT).value(), + prefixsort::PrefixSortEncoder::encodedSize(TypeKind::SMALLINT, 12) + .value(), 0, 1}, {false, - prefixsort::PrefixSortEncoder::encodedSize(TypeKind::SMALLINT).value(), + prefixsort::PrefixSortEncoder::encodedSize(TypeKind::SMALLINT, 12) + .value(), 0, 0}, - {true, 1'000'000, 0, 3}, + {true, 1'000'000, 0, 4}, {false, 1'000'000, 0, 0}, {true, 1'000'000, 1'000'000, 0}, {false, 1'000'000, 1'000'000, 0}, {true, - prefixsort::PrefixSortEncoder::encodedSize(TypeKind::SMALLINT).value() + - prefixsort::PrefixSortEncoder::encodedSize(TypeKind::INTEGER) + prefixsort::PrefixSortEncoder::encodedSize(TypeKind::SMALLINT, 12) + .value() + + prefixsort::PrefixSortEncoder::encodedSize(TypeKind::INTEGER, 12) .value(), 0, 2}, {false, - prefixsort::PrefixSortEncoder::encodedSize(TypeKind::SMALLINT).value() + - prefixsort::PrefixSortEncoder::encodedSize(TypeKind::INTEGER) + prefixsort::PrefixSortEncoder::encodedSize(TypeKind::SMALLINT, 12) + .value() + + prefixsort::PrefixSortEncoder::encodedSize(TypeKind::INTEGER, 12) .value(), 0, 0}, {true, - prefixsort::PrefixSortEncoder::encodedSize(TypeKind::SMALLINT).value() + - prefixsort::PrefixSortEncoder::encodedSize(TypeKind::INTEGER) + prefixsort::PrefixSortEncoder::encodedSize(TypeKind::SMALLINT, 12) + .value() + + prefixsort::PrefixSortEncoder::encodedSize(TypeKind::INTEGER, 12) .value(), 1'000'000, 0}, {false, - prefixsort::PrefixSortEncoder::encodedSize(TypeKind::SMALLINT).value() + - prefixsort::PrefixSortEncoder::encodedSize(TypeKind::INTEGER) + prefixsort::PrefixSortEncoder::encodedSize(TypeKind::SMALLINT, 12) + .value() + + prefixsort::PrefixSortEncoder::encodedSize(TypeKind::INTEGER, 12) .value(), 1'000'000, 0}, {true, - prefixsort::PrefixSortEncoder::encodedSize(TypeKind::SMALLINT).value() + - prefixsort::PrefixSortEncoder::encodedSize(TypeKind::BIGINT).value(), + prefixsort::PrefixSortEncoder::encodedSize(TypeKind::SMALLINT, 12) + .value() + + prefixsort::PrefixSortEncoder::encodedSize(TypeKind::BIGINT, 12) + .value(), 0, 2}, {false, - prefixsort::PrefixSortEncoder::encodedSize(TypeKind::SMALLINT).value() + - prefixsort::PrefixSortEncoder::encodedSize(TypeKind::BIGINT).value(), + prefixsort::PrefixSortEncoder::encodedSize(TypeKind::SMALLINT, 12) + .value() + + prefixsort::PrefixSortEncoder::encodedSize(TypeKind::BIGINT, 12) + .value(), 0, 0}, {true, - prefixsort::PrefixSortEncoder::encodedSize(TypeKind::SMALLINT).value() + - prefixsort::PrefixSortEncoder::encodedSize(TypeKind::BIGINT) + prefixsort::PrefixSortEncoder::encodedSize(TypeKind::SMALLINT, 12) + .value() + + prefixsort::PrefixSortEncoder::encodedSize(TypeKind::BIGINT, 12) .value() + - prefixsort::PrefixSortEncoder::encodedSize(TypeKind::INTEGER) + prefixsort::PrefixSortEncoder::encodedSize(TypeKind::INTEGER, 12) .value() - 1, 0, 2}, {false, - prefixsort::PrefixSortEncoder::encodedSize(TypeKind::SMALLINT).value() + - prefixsort::PrefixSortEncoder::encodedSize(TypeKind::BIGINT) + prefixsort::PrefixSortEncoder::encodedSize(TypeKind::SMALLINT, 12) .value() + - prefixsort::PrefixSortEncoder::encodedSize(TypeKind::INTEGER) + prefixsort::PrefixSortEncoder::encodedSize(TypeKind::BIGINT, 12) + .value() + + prefixsort::PrefixSortEncoder::encodedSize(TypeKind::INTEGER, 12) .value() - 1, 0, 0}, {true, - prefixsort::PrefixSortEncoder::encodedSize(TypeKind::SMALLINT).value() + - prefixsort::PrefixSortEncoder::encodedSize(TypeKind::BIGINT) + prefixsort::PrefixSortEncoder::encodedSize(TypeKind::SMALLINT, 12) + .value() + + prefixsort::PrefixSortEncoder::encodedSize(TypeKind::BIGINT, 12) .value() + - prefixsort::PrefixSortEncoder::encodedSize(TypeKind::INTEGER) + prefixsort::PrefixSortEncoder::encodedSize(TypeKind::INTEGER, 12) .value(), 0, 3}, {false, - prefixsort::PrefixSortEncoder::encodedSize(TypeKind::SMALLINT).value() + - prefixsort::PrefixSortEncoder::encodedSize(TypeKind::BIGINT) + prefixsort::PrefixSortEncoder::encodedSize(TypeKind::SMALLINT, 12) + .value() + + prefixsort::PrefixSortEncoder::encodedSize(TypeKind::BIGINT, 12) .value() + - prefixsort::PrefixSortEncoder::encodedSize(TypeKind::INTEGER) + prefixsort::PrefixSortEncoder::encodedSize(TypeKind::INTEGER, 12) .value(), 0, 0}}; diff --git a/velox/exec/tests/PrefixSortTest.cpp b/velox/exec/tests/PrefixSortTest.cpp index 7f4add39b7c4..90ed9db64156 100644 --- a/velox/exec/tests/PrefixSortTest.cpp +++ b/velox/exec/tests/PrefixSortTest.cpp @@ -66,7 +66,8 @@ class PrefixSortTest : public exec::test::OperatorTestBase { common::PrefixSortConfig{ 1024, // Set threshold to 0 to enable prefix-sort in small dataset. - 0}, + 0, + 12}, sortPool.get()); const auto beforeBytes = sortPool->peakBytes(); ASSERT_EQ(sortPool->peakBytes(), 0); @@ -77,7 +78,8 @@ class PrefixSortTest : public exec::test::OperatorTestBase { common::PrefixSortConfig{ 1024, // Set threshold to 0 to enable prefix-sort in small dataset. - 0}, + 0, + 12}, sortPool.get(), rows); ASSERT_GE(maxBytes, sortPool->peakBytes() - beforeBytes); @@ -151,8 +153,6 @@ const RowVectorPtr PrefixSortTest::generateExpectedResult( } TEST_F(PrefixSortTest, singleKey) { - const int columnsSize = 7; - // Vectors without nulls. const std::vector testData = { makeFlatVector({5, 4, 3, 2, 1}), @@ -172,8 +172,22 @@ TEST_F(PrefixSortTest, singleKey) { Timestamp(3, 3), Timestamp(2, 2), Timestamp(1, 1)}), - makeFlatVector({"eee", "ddd", "ccc", "bbb", "aaa"})}; - for (int i = 5; i < columnsSize; ++i) { + makeFlatVector({"eee", "ddd", "ccc", "bbb", "aaa"}), + makeFlatVector({"e", "dddddd", "bbb", "ddd", "aaa"}), + makeFlatVector( + {"ddd_not_inline", + "aaa_is_not_inline", + "aaa_is_not_inline_a", + "ddd_is_not_inline", + "aaa_is_not_inline_b"}), + makeNullableFlatVector( + {"\u7231", + "\u671B\u5E0C\u2014\u5FF5\u4FE1", + "\u671B\u5E0C", + "\u7231\u2014", + "\u4FE1 \u7231"}, + VARBINARY())}; + for (int i = 0; i < testData.size(); ++i) { const auto data = makeRowVector({testData[i]}); testPrefixSort({kAsc}, data); @@ -182,9 +196,6 @@ TEST_F(PrefixSortTest, singleKey) { } TEST_F(PrefixSortTest, singleKeyWithNulls) { - const int columnsSize = 7; - - Timestamp ts = {5, 5}; // Vectors with nulls. const std::vector testData = { makeNullableFlatVector({5, 4, std::nullopt, 2, 1}), @@ -205,9 +216,24 @@ TEST_F(PrefixSortTest, singleKeyWithNulls) { Timestamp(2, 2), Timestamp(1, 1)}), makeNullableFlatVector( - {"eee", "ddd", std::nullopt, "bbb", "aaa"})}; + {"eee", "ddd", std::nullopt, "bbb", "aaa"}), + makeNullableFlatVector( + {"ee", "aaa", std::nullopt, "d", "aaaaaaaaaaaaaa"}), + makeNullableFlatVector( + {"aaa_is_not_inline", + "aaa_is_not_inline_2", + std::nullopt, + "aaa_is_not_inline_1", + "aaaaaaaaaaaaaa"}), + makeNullableFlatVector( + {"\u7231", + "\u671B\u5E0C\u2014\u5FF5\u4FE1", + std::nullopt, + "\u7231\u2014", + "\u4FE1 \u7231"}, + VARBINARY())}; - for (int i = 5; i < columnsSize; ++i) { + for (int i = 0; i < testData.size(); ++i) { const auto data = makeRowVector({testData[i]}); testPrefixSort({kAsc}, data); @@ -232,7 +258,7 @@ TEST_F(PrefixSortTest, multipleKeys) { const auto data = makeRowVector({ makeNullableFlatVector({5, 2, std::nullopt, 2, 1}), makeNullableFlatVector( - {"eee", "ddd", std::nullopt, "bbb", "aaa"}), + {"eee", "aaa", std::nullopt, "bbb", "aaaa"}), }); testPrefixSort({kAsc, kAsc}, data); @@ -298,27 +324,33 @@ TEST_F(PrefixSortTest, checkMaxNormalizedKeySizeForMultipleKeys) { // The normalizedKeySize for BIGINT should be 8 + 1. std::vector keyTypes = {BIGINT(), BIGINT()}; std::vector compareFlags = {kAsc, kDesc}; - auto sortLayout = PrefixSortLayout::makeSortLayout(keyTypes, compareFlags, 8); + std::vector> maxStringLengths = { + std::nullopt, std::nullopt}; + auto sortLayout = PrefixSortLayout::generate( + keyTypes, compareFlags, 8, 9, maxStringLengths); ASSERT_FALSE(sortLayout.hasNormalizedKeys); - auto sortLayoutOneKey = - PrefixSortLayout::makeSortLayout(keyTypes, compareFlags, 9); + auto sortLayoutOneKey = PrefixSortLayout::generate( + keyTypes, compareFlags, 9, 9, maxStringLengths); 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); + auto sortLayoutOneKey1 = PrefixSortLayout::generate( + keyTypes, compareFlags, 17, 12, maxStringLengths); 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); + auto sortLayoutTwoKeys = PrefixSortLayout::generate( + keyTypes, compareFlags, 18, 12, maxStringLengths); ASSERT_TRUE(sortLayoutTwoKeys.hasNormalizedKeys); ASSERT_FALSE(sortLayoutTwoKeys.hasNonNormalizedKey); + ASSERT_FALSE( + sortLayoutTwoKeys.nonPrefixSortStartIndex < + sortLayoutTwoKeys.numNormalizedKeys); ASSERT_EQ(sortLayoutTwoKeys.prefixOffsets.size(), 2); ASSERT_EQ(sortLayoutTwoKeys.prefixOffsets[0], 0); ASSERT_EQ(sortLayoutTwoKeys.prefixOffsets[1], 9); @@ -346,10 +378,10 @@ TEST_F(PrefixSortTest, optimizeSortKeysOrder) { {ROW({BIGINT(), SMALLINT(), VARCHAR()}), {0, 1, 2}, {1, 0, 2}}, {ROW({TINYINT(), BIGINT(), VARCHAR(), TINYINT(), INTEGER(), VARCHAR()}), {2, 1, 0, 4, 5, 3}, - {4, 1, 2, 0, 5, 3}}, + {4, 1, 2, 5, 0, 3}}, {ROW({INTEGER(), BIGINT(), VARCHAR(), TINYINT(), INTEGER(), VARCHAR()}), {5, 4, 3, 2, 1, 0}, - {4, 0, 1, 5, 3, 2}}}; + {4, 0, 1, 5, 2, 3}}}; for (const auto& testData : testSettings) { SCOPED_TRACE(testData.debugString()); @@ -369,5 +401,48 @@ TEST_F(PrefixSortTest, optimizeSortKeysOrder) { } } } + +TEST_F(PrefixSortTest, makeSortLayoutForString) { + std::vector keyTypes = {VARCHAR(), BIGINT()}; + std::vector compareFlags = {kAsc, kDesc}; + std::vector> maxStringLengths = {9, std::nullopt}; + + auto sortLayoutOneKey = PrefixSortLayout::generate( + keyTypes, compareFlags, 24, 8, maxStringLengths); + ASSERT_TRUE(sortLayoutOneKey.hasNormalizedKeys); + ASSERT_TRUE(sortLayoutOneKey.hasNonNormalizedKey); + ASSERT_TRUE( + sortLayoutOneKey.nonPrefixSortStartIndex < + sortLayoutOneKey.numNormalizedKeys); + ASSERT_EQ(sortLayoutOneKey.encodeSizes.size(), 1); + ASSERT_EQ(sortLayoutOneKey.encodeSizes[0], 9); + + auto sortLayoutTwoCompleteKeys = PrefixSortLayout::generate( + keyTypes, compareFlags, 26, 9, maxStringLengths); + ASSERT_TRUE(sortLayoutTwoCompleteKeys.hasNormalizedKeys); + ASSERT_FALSE(sortLayoutTwoCompleteKeys.hasNonNormalizedKey); + ASSERT_TRUE( + sortLayoutTwoCompleteKeys.nonPrefixSortStartIndex == + sortLayoutTwoCompleteKeys.numNormalizedKeys); + ASSERT_EQ(sortLayoutTwoCompleteKeys.encodeSizes.size(), 2); + ASSERT_EQ(sortLayoutTwoCompleteKeys.encodeSizes[0], 10); + ASSERT_EQ(sortLayoutTwoCompleteKeys.encodeSizes[1], 9); + + // The last key type is VARBINARY, indicating that only partial data is + // encoded in the prefix. + std::vector keyTypes1 = {BIGINT(), VARBINARY()}; + maxStringLengths = {std::nullopt, 9}; + auto sortLayoutTwoKeys = PrefixSortLayout::generate( + keyTypes1, compareFlags, 26, 8, maxStringLengths); + ASSERT_TRUE(sortLayoutTwoKeys.hasNormalizedKeys); + ASSERT_FALSE(sortLayoutTwoKeys.hasNonNormalizedKey); + ASSERT_TRUE( + sortLayoutTwoKeys.nonPrefixSortStartIndex < + sortLayoutTwoKeys.numNormalizedKeys); + ASSERT_EQ(sortLayoutTwoKeys.encodeSizes.size(), 2); + ASSERT_EQ(sortLayoutTwoKeys.encodeSizes[0], 9); + ASSERT_EQ(sortLayoutTwoKeys.encodeSizes[1], 9); +} + } // namespace } // namespace facebook::velox::exec::prefixsort::test diff --git a/velox/exec/tests/RowContainerTest.cpp b/velox/exec/tests/RowContainerTest.cpp index 4197ea5c0cbb..a2040721f664 100644 --- a/velox/exec/tests/RowContainerTest.cpp +++ b/velox/exec/tests/RowContainerTest.cpp @@ -2597,4 +2597,44 @@ TEST_F(RowContainerTest, rowColumnStats) { EXPECT_EQ(stats.nullCount(), 4); EXPECT_EQ(stats.numCells(), 10); } + +TEST_F(RowContainerTest, storeAndCollectColumnStats) { + const uint64_t kNumRows = 1000; + auto rowVector = makeRowVector({ + makeFlatVector( + kNumRows, [](auto row) { return row % 5; }, nullEvery(7)), + makeFlatVector( + kNumRows, + [](auto row) { return fmt::format("abcdefg123_{}", row); }, + nullEvery(7)), + }); + + auto rowContainer = makeRowContainer({BIGINT(), VARCHAR()}, {}, false); + std::vector rows; + rows.reserve(kNumRows); + + SelectivityVector allRows(kNumRows); + for (size_t i = 0; i < kNumRows; i++) { + auto row = rowContainer->newRow(); + rows.push_back(row); + } + for (int i = 0; i < rowContainer->columnTypes().size(); ++i) { + DecodedVector decoded(*rowVector->childAt(i), allRows); + rowContainer->store(decoded, folly::Range(rows.data(), kNumRows), i); + } + + ASSERT_EQ(rowContainer->numRows(), kNumRows); + for (int i = 0; i < rowContainer->columnTypes().size(); ++i) { + const auto stats = rowContainer->columnStats(i).value(); + EXPECT_EQ(stats.nonNullCount(), 857); + EXPECT_EQ(stats.nullCount(), 143); + EXPECT_EQ(stats.numCells(), kNumRows); + if (rowVector->childAt(i)->typeKind() == TypeKind::VARCHAR) { + EXPECT_EQ(stats.maxBytes(), 14); + EXPECT_EQ(stats.minBytes(), 12); + EXPECT_EQ(stats.sumBytes(), 11905); + EXPECT_EQ(stats.avgBytes(), 13); + } + } +} } // namespace facebook::velox::exec::test diff --git a/velox/exec/tests/SortBufferTest.cpp b/velox/exec/tests/SortBufferTest.cpp index b4feb8d667a4..93bf32edf4df 100644 --- a/velox/exec/tests/SortBufferTest.cpp +++ b/velox/exec/tests/SortBufferTest.cpp @@ -66,7 +66,13 @@ class SortBufferTest : public OperatorTestBase, OperatorTestBase::TearDown(); } - common::SpillConfig getSpillConfig(const std::string& spillDir) const { + common::SpillConfig getSpillConfig( + const std::string& spillDir, + bool enableSpillPrefixSort = true) const { + std::optional spillPrefixSortConfig = + enableSpillPrefixSort + ? std::optional(prefixSortConfig_) + : std::nullopt; return common::SpillConfig( [spillDir]() -> const std::string& { return spillDir; }, [&](uint64_t) {}, @@ -83,14 +89,15 @@ class SortBufferTest : public OperatorTestBase, 0, 0, "none", - spillPrefixSortConfig_); + spillPrefixSortConfig); } const bool enableSpillPrefixSort_{GetParam()}; const velox::common::PrefixSortConfig prefixSortConfig_ = velox::common::PrefixSortConfig{ std::numeric_limits::max(), - GetParam() ? 8 : std::numeric_limits::max()}; + GetParam() ? 8 : std::numeric_limits::max(), + 12}; const std::optional spillPrefixSortConfig_ = enableSpillPrefixSort_ ? std::optional(prefixSortConfig_) @@ -103,13 +110,6 @@ class SortBufferTest : public OperatorTestBase, {"c3", REAL()}, {"c4", DOUBLE()}, {"c5", VARCHAR()}}); - const RowTypePtr nonPrefixSortInputType_ = ROW( - {{"c0", VARCHAR()}, - {"c1", VARCHAR()}, - {"c2", VARCHAR()}, - {"c3", VARCHAR()}, - {"c4", VARCHAR()}, - {"c5", VARCHAR()}}); // Specifies the sort columns ["c4", "c1"]. std::vector sortColumnIndices_{4, 1}; std::vector sortCompareFlags_{ @@ -738,12 +738,10 @@ DEBUG_ONLY_TEST_P(SortBufferTest, reserveMemorySort) { SCOPED_TRACE(fmt::format( "usePrefixSort: {}, spillEnabled: {}, ", usePrefixSort, spillEnabled)); auto spillDirectory = exec::test::TempDirectoryPath::create(); - auto spillConfig = getSpillConfig(spillDirectory->getPath()); + auto spillConfig = getSpillConfig(spillDirectory->getPath(), usePrefixSort); folly::Synchronized spillStats; - const RowTypePtr inputType = - usePrefixSort ? inputType_ : nonPrefixSortInputType_; auto sortBuffer = std::make_unique( - inputType, + inputType_, sortColumnIndices_, sortCompareFlags_, pool_.get(), @@ -757,7 +755,7 @@ DEBUG_ONLY_TEST_P(SortBufferTest, reserveMemorySort) { VectorFuzzer fuzzer({.vectorSize = 100}, spillSource.get()); TestScopedSpillInjection scopedSpillInjection(0); - sortBuffer->addInput(fuzzer.fuzzRow(inputType)); + sortBuffer->addInput(fuzzer.fuzzRow(inputType_)); std::atomic_bool hasReserveMemory = false; // Reserve memory for sort. diff --git a/velox/exec/tests/WindowTest.cpp b/velox/exec/tests/WindowTest.cpp index a89490283804..ffaac48b1493 100644 --- a/velox/exec/tests/WindowTest.cpp +++ b/velox/exec/tests/WindowTest.cpp @@ -659,7 +659,7 @@ DEBUG_ONLY_TEST_F(WindowTest, reserveMemorySort) { const auto plan = usePrefixSort ? prefixSortPlan : nonPrefixSortPlan; velox::common::PrefixSortConfig prefixSortConfig = velox::common::PrefixSortConfig{ - std::numeric_limits::max(), 130}; + std::numeric_limits::max(), 130, 12}; auto sortWindowBuild = std::make_unique( plan, pool_.get(), diff --git a/velox/exec/tests/utils/LocalRunnerTestBase.cpp b/velox/exec/tests/utils/LocalRunnerTestBase.cpp index 2c23b50a3246..123a4f9e8aa8 100644 --- a/velox/exec/tests/utils/LocalRunnerTestBase.cpp +++ b/velox/exec/tests/utils/LocalRunnerTestBase.cpp @@ -59,7 +59,7 @@ void LocalRunnerTestBase::ensureTestData() { void LocalRunnerTestBase::makeSchema() { auto schemaQueryCtx = makeQueryCtx("schema", rootPool_.get()); common::SpillConfig spillConfig; - common::PrefixSortConfig prefixSortConfig(100, 130); + common::PrefixSortConfig prefixSortConfig(100, 130, 12); auto leafPool = schemaQueryCtx->pool()->addLeafChild("schemaReader"); auto connectorQueryCtx = std::make_shared( leafPool.get(),