diff --git a/bolt/common/memory/RawVector.cpp b/bolt/common/memory/RawVector.cpp index b165229c3..8c1116aab 100644 --- a/bolt/common/memory/RawVector.cpp +++ b/bolt/common/memory/RawVector.cpp @@ -43,13 +43,14 @@ bool initializeIota() { } } // namespace -const int32_t* iota(int32_t size, raw_vector& storage) { - if (iotaData.size() < size) { +const int32_t* +iota(int32_t size, raw_vector& storage, int32_t offset) { + if (iotaData.size() < offset + size) { storage.resize(size); - std::iota(&storage[0], &storage[storage.size()], 0); + std::iota(storage.begin(), storage.end(), offset); return storage.data(); } - return iotaData.data(); + return iotaData.data() + offset; } static bool FB_ANONYMOUS_VARIABLE(g_iotaConstants) = initializeIota(); diff --git a/bolt/common/memory/RawVector.h b/bolt/common/memory/RawVector.h index 26b79bdc7..ec76f7120 100644 --- a/bolt/common/memory/RawVector.h +++ b/bolt/common/memory/RawVector.h @@ -266,11 +266,12 @@ class raw_vector { }; // Returns a pointer to 'size' int32_t's with consecutive values -// starting at 0. There are at least kPadding / sizeof(int32_t) values +// starting at 'offset'. There are at least kPadding / sizeof(int32_t) values // past 'size', so that it is safe to access the returned pointer at maximum // SIMD width. Typically returns preallocated memory but if this is // not large enough,resizes and initializes 'storage' to the requested // size and returns storage.data(). -const int32_t* iota(int32_t size, raw_vector& storage); +const int32_t* +iota(int32_t size, raw_vector& storage, int32_t offset = 0); } // namespace bytedance::bolt diff --git a/bolt/common/memory/tests/RawVectorTest.cpp b/bolt/common/memory/tests/RawVectorTest.cpp index b7f5aae33..6481c3623 100644 --- a/bolt/common/memory/tests/RawVectorTest.cpp +++ b/bolt/common/memory/tests/RawVectorTest.cpp @@ -134,10 +134,15 @@ TEST_P(RawVectorTest, iota) { raw_vector storage(pool_.get()); // Small sizes are preallocated. EXPECT_EQ(11, iota(12, storage)[11]); + EXPECT_EQ(10, iota(12, storage, 10)[0]); + EXPECT_EQ(21, iota(12, storage, 10)[11]); EXPECT_TRUE(storage.empty()); EXPECT_EQ(110000, iota(110001, storage)[110000]); // Larger sizes are allocated in 'storage'. EXPECT_FALSE(storage.empty()); + EXPECT_EQ(20000, iota(3, storage, 20000)[0]); + EXPECT_EQ(20002, iota(3, storage, 20000)[2]); + EXPECT_EQ(3, storage.size()); } TEST_P(RawVectorTest, iterator) { diff --git a/bolt/connectors/hive/HiveDataSource.cpp b/bolt/connectors/hive/HiveDataSource.cpp index 1929c0c41..ae9c16949 100644 --- a/bolt/connectors/hive/HiveDataSource.cpp +++ b/bolt/connectors/hive/HiveDataSource.cpp @@ -157,10 +157,14 @@ HiveDataSource::HiveDataSource( if (remainingFilter) { remainingFilterExprSet_ = expressionEvaluator_->compile(remainingFilter); auto& remainingFilterExpr = remainingFilterExprSet_->expr(0); - folly::F14FastSet columnNames( - readerRowNames.begin(), readerRowNames.end()); + folly::F14FastMap columnNames; + for (column_index_t i = 0; i < readerRowNames.size(); ++i) { + columnNames[readerRowNames[i]] = i; + } for (auto& input : remainingFilterExpr->distinctFields()) { - if (columnNames.count(input->field()) > 0) { + auto it = columnNames.find(input->field()); + if (it != columnNames.end()) { + multiReferencedFields_.push_back(it->second); continue; } // Remaining filter may reference columns that are not used otherwise, @@ -849,9 +853,16 @@ int64_t HiveDataSource::estimatedRowSize() { } vector_size_t HiveDataSource::evaluateRemainingFilter(RowVectorPtr& rowVector) { - auto filterStartMicros = getCurrentTimeMicro(); filterRows_.resize(output_->size()); + for (auto fieldIndex : multiReferencedFields_) { + LazyVector::ensureLoadedRows( + rowVector->childAt(fieldIndex), + filterRows_, + filterLazyDecoded_, + filterLazyBaseRows_); + } + auto filterStartMicros = getCurrentTimeMicro(); expressionEvaluator_->evaluate( remainingFilterExprSet_.get(), filterRows_, diff --git a/bolt/connectors/hive/HiveDataSource.h b/bolt/connectors/hive/HiveDataSource.h index 4757d5a30..367f43777 100644 --- a/bolt/connectors/hive/HiveDataSource.h +++ b/bolt/connectors/hive/HiveDataSource.h @@ -45,6 +45,8 @@ #include "bolt/dwio/common/Statistics.h" #include "bolt/exec/OperatorUtils.h" #include "bolt/expression/Expr.h" +#include "bolt/vector/DecodedVector.h" +#include "bolt/vector/LazyVector.h" namespace bytedance::bolt::connector::hive { class HiveConfig; @@ -210,9 +212,15 @@ class HiveDataSource : public DataSource { std::atomic totalRemainingFilterTime_{0}; uint64_t completedRows_ = 0; + // Field indices referenced in both remaining filter and output type. These + // columns need to be materialized eagerly to avoid missing values in output. + std::vector multiReferencedFields_; + // Reusable memory for remaining filter evaluation. VectorPtr filterResult_; SelectivityVector filterRows_; + DecodedVector filterLazyDecoded_; + SelectivityVector filterLazyBaseRows_; exec::FilterEvalCtx filterEvalCtx_; RowVectorPtr emptyResult_; diff --git a/bolt/connectors/hive/PaimonSplitReader.cpp b/bolt/connectors/hive/PaimonSplitReader.cpp index a9e091e3f..19de8ac90 100644 --- a/bolt/connectors/hive/PaimonSplitReader.cpp +++ b/bolt/connectors/hive/PaimonSplitReader.cpp @@ -113,6 +113,7 @@ PaimonRowIteratorPtr PaimonSplitReader::getIterator(SplitReader* rowReader) { if (rowReader->next(paimon::kMAX_BATCH_SIZE, result)) { RowVectorPtr resultAsRowVect = std::static_pointer_cast(result); + resultAsRowVect->loadedVector(); auto primaryKeys = projectVector(resultAsRowVect, primaryKeyIndices_); auto sequenceFields = diff --git a/bolt/dwio/common/ColumnVisitors.h b/bolt/dwio/common/ColumnVisitors.h index c7a02643a..80b0d4e82 100644 --- a/bolt/dwio/common/ColumnVisitors.h +++ b/bolt/dwio/common/ColumnVisitors.h @@ -168,6 +168,10 @@ class ColumnVisitor { using DataType = T; static constexpr bool dense = isDense; static constexpr bool kHasBulkPath = hasBulkPath; + static constexpr bool kHasFilter = + !std::is_same_v; + static constexpr bool kHasHook = + !std::is_same_v; ColumnVisitor( TFilter& filter, SelectiveColumnReader* reader, @@ -433,6 +437,10 @@ class ColumnVisitor { numValuesBias_ = bias; } + int32_t numValuesBias() const { + return numValuesBias_; + } + void setNumValues(int32_t size) { reader_->setNumValues(numValuesBias_ + size); if constexpr (!std::is_same_v) { @@ -528,7 +536,7 @@ template < inline void ColumnVisitor::addResult( T value) { - values_.addValue(rowIndex_, value); + values_.addValue(rowIndex_ + numValuesBias_, value); } template < @@ -539,7 +547,7 @@ template < bool hasBulkPath> inline void ColumnVisitor::addNull() { - values_.template addNull(rowIndex_); + values_.template addNull(rowIndex_ + numValuesBias_); } template < @@ -857,7 +865,10 @@ class DictionaryColumnVisitor translateByDict(input, numInput, values); super::values_.hook().addValues( scatter ? scatterRows + super::rowIndex_ - : bolt::iota(super::numRows_, super::innerNonNullRows()) + + : bolt::iota( + super::numRows_, + super::innerNonNullRows(), + super::numValuesBias_) + super::rowIndex_, values, numInput, @@ -1139,6 +1150,12 @@ class DictionaryColumnVisitor return state_.dictionary.numValues; } + static constexpr bool hasFilter() { + // Dictionary values cannot be null. + return !std::is_same_v && + !std::is_same_v; + } + uint8_t* filterCache() const { return state_.filterCache; } @@ -1214,8 +1231,13 @@ class StringDictionaryColumnVisitor } vector_size_t previous = isDense && TFilter::deterministic ? 0 : super::currentRow(); - if constexpr (std::is_same_v) { - super::filterPassed(index); + if constexpr (!DictSuper::hasFilter()) { + if constexpr (super::kHasHook) { + super::values_.addValue( + super::rowIndex_ + super::numValuesBias_, valueInDictionary(index)); + } else { + super::filterPassed(index); + } } else { // check the dictionary cache if (TFilter::deterministic && @@ -1227,7 +1249,7 @@ class StringDictionaryColumnVisitor super::filterFailed(); } else { if (bolt::common::applyFilter( - super::filter_, valueInDictionary(value, inStrideDict))) { + super::filter_, valueInDictionary(index))) { super::filterPassed(index); if constexpr (TFilter::deterministic) { DictSuper::filterCache()[index] = FilterResult::kSuccess; @@ -1270,25 +1292,31 @@ class StringDictionaryColumnVisitor int32_t& numValues) { DCHECK(input == values + numValues); setByInDict(values + numValues, numInput); - if (!hasFilter) { + if constexpr (!DictSuper::hasFilter()) { if (hasHook) { for (auto i = 0; i < numInput; ++i) { - auto value = input[i]; super::values_.addValue( scatterRows ? scatterRows[super::rowIndex_ + i] - : super::rowIndex_ + i, - value); + : super::rowIndex_ + super::numValuesBias_ + i, + valueInDictionary(input[i])); } } - DCHECK_EQ(input, values + numValues); - if (scatter) { + if constexpr (std::is_same_v) { + auto* begin = (scatter ? scatterRows : super::rows_) + super::rowIndex_; + std::copy(begin, begin + numInput, filterHits + numValues); + numValues += numInput; + } else if (scatter) { dwio::common::scatterDense( input, scatterRows + super::rowIndex_, numInput, values); + numValues = scatterRows[super::rowIndex_ + numInput - 1] + 1; + } else { + DCHECK_EQ(input, values + numValues); + numValues += numInput; } - numValues = scatter ? scatterRows[super::rowIndex_ + numInput - 1] + 1 - : numValues + numInput; super::rowIndex_ += numInput; return; + } else { + static_assert(hasFilter); } constexpr bool filterOnly = std::is_same_v; @@ -1325,16 +1353,7 @@ class StringDictionaryColumnVisitor while (bits) { int index = bits::getAndClearLastSetBit(bits); int32_t value = input[i + index]; - bool result; - if (value >= DictSuper::dictionarySize()) { - result = applyFilter( - super::filter_, - valueInDictionary(value - DictSuper::dictionarySize(), true)); - } else { - result = - applyFilter(super::filter_, valueInDictionary(value, false)); - } - if (result) { + if (applyFilter(super::filter_, valueInDictionary(value))) { DictSuper::filterCache()[value] = FilterResult::kSuccess; passed |= 1 << index; } else { @@ -1414,67 +1433,17 @@ class StringDictionaryColumnVisitor } } - folly::StringPiece valueInDictionary(int64_t index, bool inStrideDict) { - if (inStrideDict) { - return folly::StringPiece(reinterpret_cast( - DictSuper::state_.dictionary2.values)[index]); + folly::StringPiece valueInDictionary(int64_t index) { + auto stripeDictSize = DictSuper::state_.dictionary.numValues; + if (index < stripeDictSize) { + return reinterpret_cast( + DictSuper::state_.dictionary.values)[index]; } - return folly::StringPiece(reinterpret_cast( - DictSuper::state_.dictionary.values)[index]); + return reinterpret_cast( + DictSuper::state_.dictionary2.values)[index - stripeDictSize]; } }; -class ExtractStringDictionaryToGenericHook { - public: - static constexpr bool kSkipNulls = true; - using HookType = ValueHook; - - ExtractStringDictionaryToGenericHook( - ValueHook* hook, - RowSet rows, - RawScanState state) - - : hook_(hook), rows_(rows), state_(state) {} - - bool acceptsNulls() { - return hook_->acceptsNulls(); - } - - template - void addNull(vector_size_t rowIndex) { - hook_->addNull(rowIndex); - } - - void addValue(vector_size_t rowIndex, int32_t value) { - // We take the string from the stripe or stride dictionary - // according to the index. Stride dictionary indices are offset up - // by the stripe dict size. - if (value < dictionarySize()) { - auto view = folly::StringPiece( - reinterpret_cast(state_.dictionary.values)[value]); - hook_->addValue(rowIndex, &view); - } else { - BOLT_DCHECK(state_.inDictionary); - auto view = folly::StringPiece(reinterpret_cast( - state_.dictionary2.values)[value - dictionarySize()]); - hook_->addValue(rowIndex, &view); - } - } - - ValueHook& hook() { - return *hook_; - } - - private: - int32_t dictionarySize() const { - return state_.dictionary.numValues; - } - - ValueHook* const hook_; - RowSet const rows_; - RawScanState state_; -}; - template class DirectRleColumnVisitor : public ColumnVisitor { diff --git a/bolt/dwio/common/DirectDecoder.h b/bolt/dwio/common/DirectDecoder.h index 2a6517a67..9a420bc5d 100644 --- a/bolt/dwio/common/DirectDecoder.h +++ b/bolt/dwio/common/DirectDecoder.h @@ -221,6 +221,11 @@ class DirectDecoder : public IntDecoder { return; } } + if (hasHook && visitor.numValuesBias() > 0) { + for (auto& row : *outerVector) { + row += visitor.numValuesBias(); + } + } if (super::useVInts) { if (Visitor::dense) { super::bulkRead(numNonNull, data); @@ -271,7 +276,10 @@ class DirectDecoder : public IntDecoder { rowsAsRange, 0, rowsAsRange.size(), - hasHook ? bolt::iota(numRows, visitor.innerNonNullRows()) + hasHook ? bolt::iota( + numRows, + visitor.innerNonNullRows(), + visitor.numValuesBias()) : nullptr, visitor.rawValues(numRows), hasFilter ? visitor.outputRows(numRows) : nullptr, @@ -281,7 +289,11 @@ class DirectDecoder : public IntDecoder { } else { dwio::common::fixedWidthScan( rowsAsRange, - hasHook ? bolt::iota(numRows, visitor.innerNonNullRows()) : nullptr, + hasHook ? bolt::iota( + numRows, + visitor.innerNonNullRows(), + visitor.numValuesBias()) + : nullptr, visitor.rawValues(numRows), hasFilter ? visitor.outputRows(numRows) : nullptr, numValues, diff --git a/bolt/dwio/common/SelectiveStructColumnReader.cpp b/bolt/dwio/common/SelectiveStructColumnReader.cpp index e5a57407f..57a7a0955 100644 --- a/bolt/dwio/common/SelectiveStructColumnReader.cpp +++ b/bolt/dwio/common/SelectiveStructColumnReader.cpp @@ -187,8 +187,9 @@ void SelectiveStructColumnReaderBase::read( } auto fieldIndex = childSpec->subscript(); auto reader = children_.at(fieldIndex); - if (reader->isTopLevel() && childSpec->projectOut() && - !childSpec->hasFilter() && !childSpec->extractValues()) { + if (!shouldReadChildrenEagerly() && reader->isTopLevel() && + childSpec->projectOut() && !childSpec->hasFilter() && + !childSpec->extractValues()) { // Will make a LazyVector. continue; } @@ -432,8 +433,8 @@ void SelectiveStructColumnReaderBase::getValues( setNullField(rows.size(), childResult, childType, resultRow->pool()); continue; } - if (childSpec->extractValues() || childSpec->hasFilter() || - !children_[index]->isTopLevel()) { + if (shouldReadChildrenEagerly() || childSpec->extractValues() || + childSpec->hasFilter() || !children_[index]->isTopLevel()) { children_[index]->getValues(rows, &childResult); continue; } @@ -467,8 +468,8 @@ void SelectiveStructColumnReaderBase::getValues( auto dcKeys = fileType_->getDcKeys(); const auto it = dcKeys.find(scanSpec_->fieldName()); if (it != dcKeys.end()) { - auto oldMap = resultRow->childAt(0)->as(); - auto oldRow = resultRow->childAt(1)->as(); + auto oldMap = resultRow->childAt(0)->loadedVector()->as(); + auto oldRow = resultRow->childAt(1)->loadedVector()->as(); const auto& keys = it->second; BOLT_CHECK_EQ(oldRow->childrenSize(), keys.size()); diff --git a/bolt/dwio/common/SelectiveStructColumnReader.h b/bolt/dwio/common/SelectiveStructColumnReader.h index b41907c76..32ce8e1b9 100644 --- a/bolt/dwio/common/SelectiveStructColumnReader.h +++ b/bolt/dwio/common/SelectiveStructColumnReader.h @@ -157,6 +157,10 @@ class SelectiveStructColumnReaderBase : public SelectiveColumnReader { bool isChildMissing(const bolt::common::ScanSpec& childSpec) const; + virtual bool shouldReadChildrenEagerly() const { + return false; + } + const dwio::common::ColumnReaderOptions columnReaderOptions_; const std::shared_ptr requestedType_; diff --git a/bolt/dwio/dwrf/reader/SelectiveStringDictionaryColumnReader.cpp b/bolt/dwio/dwrf/reader/SelectiveStringDictionaryColumnReader.cpp index c59c03717..692c1ad39 100644 --- a/bolt/dwio/dwrf/reader/SelectiveStringDictionaryColumnReader.cpp +++ b/bolt/dwio/dwrf/reader/SelectiveStringDictionaryColumnReader.cpp @@ -260,14 +260,12 @@ void SelectiveStringDictionaryColumnReader::read( readHelper( &alwaysTrue(), rows, - ExtractStringDictionaryToGenericHook( - scanSpec_->valueHook(), rows, scanState_.rawState)); + dwio::common::ExtractToGenericHook(scanSpec_->valueHook())); } else { readHelper( &alwaysTrue(), rows, - ExtractStringDictionaryToGenericHook( - scanSpec_->valueHook(), rows, scanState_.rawState)); + dwio::common::ExtractToGenericHook(scanSpec_->valueHook())); } } else { if (isDense) { diff --git a/bolt/dwio/dwrf/test/SelectiveStringDictionaryColumnReaderTest.cpp b/bolt/dwio/dwrf/test/SelectiveStringDictionaryColumnReaderTest.cpp index 5cba5223e..51ed7e8c0 100755 --- a/bolt/dwio/dwrf/test/SelectiveStringDictionaryColumnReaderTest.cpp +++ b/bolt/dwio/dwrf/test/SelectiveStringDictionaryColumnReaderTest.cpp @@ -63,6 +63,24 @@ class SelectiveStringDictionaryColumnReaderTest : public ::testing::Test { } }; +class CollectStringHook : public ValueHook { + public: + void addValue(vector_size_t row, const void* value) override { + if (row >= values_.size()) { + values_.resize(row + 1); + } + values_[row] = + std::string(*reinterpret_cast(value)); + } + + const std::vector& values() const { + return values_; + } + + private: + std::vector values_; +}; + // Helper to build a TypeWithId for a single string column "myString". std::shared_ptr makeFileType() { auto rowType = HiveTypeParser().parse("struct"); @@ -82,6 +100,76 @@ std::unique_ptr makeReader( } // namespace +TEST_F( + SelectiveStringDictionaryColumnReaderTest, + ValueHookReceivesDecodedDictionaryStrings) { + MockStripeStreams streams; + memory::AllocationPool pool{&streams.getMemoryPool()}; + StreamLabels labels{pool}; + ColumnReaderStatistics columnStats; + DwrfParams params(streams, labels, columnStats); + + proto::ColumnEncoding dictEncoding; + dictEncoding.set_kind(proto::ColumnEncoding_Kind_DICTIONARY); + dictEncoding.set_dictionarysize(3); + EXPECT_CALL(streams, getEncodingProxy(_)) + .WillRepeatedly(Return(&dictEncoding)); + + EXPECT_CALL(streams, getStreamProxy(_, proto::Stream_Kind_PRESENT, false)) + .WillRepeatedly(Return(nullptr)); + EXPECT_CALL(streams, getStreamProxy(_, proto::Stream_Kind_ROW_INDEX, false)) + .WillRepeatedly(Return(nullptr)); + EXPECT_CALL( + streams, getStreamProxy(1, proto::Stream_Kind_IN_DICTIONARY, false)) + .WillRepeatedly(Return(nullptr)); + + char data[16]; + size_t dataLen = 1; + data[0] = 0x9c; // -100 literal run. + const std::vector indices = {0, 1, 2, 1}; + for (auto index : indices) { + dataLen = writeVuLong(data, dataLen, index); + } + EXPECT_CALL(streams, getStreamProxy(1, proto::Stream_Kind_DATA, true)) + .WillRepeatedly(Return(new SeekableArrayInputStream(data, dataLen))); + + const std::string dict = "alphabetagamma"; + EXPECT_CALL( + streams, getStreamProxy(1, proto::Stream_Kind_DICTIONARY_DATA, false)) + .WillRepeatedly( + Return(new SeekableArrayInputStream(dict.data(), dict.size()))); + + char lengths[16]; + size_t lengthsLen = 1; + lengths[0] = 0xfd; // -3 literal run. + for (auto length : {5, 4, 5}) { + lengthsLen = writeVuLong(lengths, lengthsLen, length); + } + EXPECT_CALL(streams, getStreamProxy(1, proto::Stream_Kind_LENGTH, false)) + .WillRepeatedly( + Return(new SeekableArrayInputStream(lengths, lengthsLen))); + + StaticStrideIndexProvider provider(0); + EXPECT_CALL(streams, getStrideIndexProviderProxy()) + .WillRepeatedly(Return(&provider)); + + auto fileType = makeFileType(); + common::ScanSpec scanSpec("myString"); + scanSpec.setProjectOut(true); + scanSpec.setExtractValues(true); + CollectStringHook hook; + scanSpec.setValueHook(&hook); + + auto reader = makeReader(fileType, params, scanSpec); + std::vector rowNumbers(indices.size()); + std::iota(rowNumbers.begin(), rowNumbers.end(), 0); + RowSet rowSet(rowNumbers.data(), rowNumbers.data() + rowNumbers.size()); + + reader->read(0, rowSet, nullptr); + + EXPECT_THAT(hook.values(), ElementsAre("alpha", "beta", "gamma", "beta")); +} + // 1) Happy path: stride dictionary streams (data + length) are present and // loadStrideDictionary populates scanState_.dictionary2. TEST_F( diff --git a/bolt/dwio/paimon/deletionvectors/tests/DeletionFileReaderTest.cpp b/bolt/dwio/paimon/deletionvectors/tests/DeletionFileReaderTest.cpp index 9cfcd0055..d78d0ce9d 100644 --- a/bolt/dwio/paimon/deletionvectors/tests/DeletionFileReaderTest.cpp +++ b/bolt/dwio/paimon/deletionvectors/tests/DeletionFileReaderTest.cpp @@ -78,6 +78,7 @@ class DeleteionFileReaderTest : public parquet::ParquetTestBase { while (true) { auto part = rowReader->next(1024, output); if (part > 0) { + output->loadedVector(); data.emplace_back(output); } else { break; @@ -121,6 +122,7 @@ class DeleteionFileReaderTest : public parquet::ParquetTestBase { mutation.deletedRows = deletionVector->as(); auto part = reader->next(batchSize, result, &mutation); if (part > 0) { + result->loadedVector(); assertEqualVectorPart(expected, result, total); total += result->size(); } else { diff --git a/bolt/dwio/paimon/reader/tests/PaimonReaderMetadataFieldTest.cpp b/bolt/dwio/paimon/reader/tests/PaimonReaderMetadataFieldTest.cpp index 9d75577a0..21886a7ee 100644 --- a/bolt/dwio/paimon/reader/tests/PaimonReaderMetadataFieldTest.cpp +++ b/bolt/dwio/paimon/reader/tests/PaimonReaderMetadataFieldTest.cpp @@ -185,6 +185,7 @@ class PaimonReaderMetadataFieldTest // accumulate results into a single vector std::vector results; while (auto result = readTask->next()) { + result->loadedVector(); results.push_back(result); } @@ -874,6 +875,7 @@ TEST_F(PaimonReaderMetadataFieldTest, testPaimonFilePathColumnMultiFile) { readTask->noMoreSplits(scanNodeId); std::vector results; while (auto result = readTask->next()) { + result->loadedVector(); results.push_back(result); } diff --git a/bolt/dwio/parquet/reader/PageReader.h b/bolt/dwio/parquet/reader/PageReader.h index 51038529e..7212e7176 100644 --- a/bolt/dwio/parquet/reader/PageReader.h +++ b/bolt/dwio/parquet/reader/PageReader.h @@ -457,14 +457,17 @@ class PageReader { // Returns the number of passed rows/values gathered by // 'reader'. Only numRows() is set for a filter-only case, only // numValues() is set for a non-filtered case. - template - static int32_t numRowsInReader( - const dwio::common::SelectiveColumnReader& reader) { + template + static int32_t numValuesRead( + const dwio::common::SelectiveColumnReader& reader, + int32_t numPageRowsRead) { + if (hasHook) { + return numPageRowsRead; + } if (hasFilter) { return reader.numRows(); - } else { - return reader.numValues(); } + return reader.numValues(); } void preloadPageRepDefs(const bool keepRepDefRawData); @@ -658,6 +661,9 @@ void PageReader::readWithVisitor(Visitor& visitor) { !std::is_same_v; constexpr bool filterOnly = std::is_same_v; + constexpr bool hasHook = Visitor::kHasHook; + static_assert( + !(hasFilter && hasHook), "hasFilter and hasHook cannot both be true"); bool mayProduceNulls = !filterOnly && visitor.allowNulls(); auto rows = visitor.rows(); auto numRows = visitor.numRows(); @@ -665,11 +671,13 @@ void PageReader::readWithVisitor(Visitor& visitor) { startVisit(folly::Range(rows, numRows)); rowsCopy_ = &visitor.rowsCopy(); folly::Range pageRows; + int32_t numPageRowsRead = 0; const uint64_t* nulls = nullptr; bool isMultiPage = false; while (rowsForPage(reader, hasFilter, mayProduceNulls, pageRows, nulls)) { bool nullsFromFastPath = false; - int32_t numValuesBeforePage = numRowsInReader(reader); + const int32_t numValuesBeforePage = + numValuesRead(reader, numPageRowsRead); visitor.setNumValuesBias(numValuesBeforePage); visitor.setRows(pageRows); { @@ -678,9 +686,12 @@ void PageReader::readWithVisitor(Visitor& visitor) { callDecoder(nulls, nullsFromFastPath, visitor); } const auto numValuesFromPage = - numRowsInReader(reader) - numValuesBeforePage; + numValuesRead(reader, numPageRowsRead) - + numValuesBeforePage; + const auto processedPageValues = + hasHook ? static_cast(pageRows.size()) : numValuesFromPage; if (currentPageNumValues_ >= 0) { - currentPageNumValues_ += numValuesFromPage; + currentPageNumValues_ += processedPageValues; } if (encoding_ == thrift::Encoding::DELTA_BINARY_PACKED && deltaBpDecoder_->validValuesCount() == 0) { @@ -726,6 +737,7 @@ void PageReader::readWithVisitor(Visitor& visitor) { if (hasFilter && rowNumberBias_) { reader.offsetOutputRows(numValuesBeforePage, rowNumberBias_); } + numPageRowsRead += pageRows.size(); if (currentPageNumValues_ >= 0 && firstUnvisited_ == rowOfPage_ + numRowsInPage_) { if (currentPageNumValues_ == 0 && FOLLY_LIKELY(statis_ != nullptr)) { diff --git a/bolt/dwio/parquet/reader/ParquetReader.cpp b/bolt/dwio/parquet/reader/ParquetReader.cpp index 05fe426e5..027157689 100644 --- a/bolt/dwio/parquet/reader/ParquetReader.cpp +++ b/bolt/dwio/parquet/reader/ParquetReader.cpp @@ -1557,6 +1557,7 @@ class ParquetRowReader::Impl { params, *options_.getScanSpec(), pool_); + columnReader_->setIsTopLevel(); extractRequiredColumnNodes(); diff --git a/bolt/dwio/parquet/reader/RleBpDataDecoder.h b/bolt/dwio/parquet/reader/RleBpDataDecoder.h index 26471a031..17fa69073 100644 --- a/bolt/dwio/parquet/reader/RleBpDataDecoder.h +++ b/bolt/dwio/parquet/reader/RleBpDataDecoder.h @@ -133,6 +133,11 @@ class RleBpDataDecoder : public bytedance::bolt::parquet::RleBpDecoder { visitor.setAllNull(hasFilter ? 0 : numRows); return; } + if (hasHook && visitor.numValuesBias() > 0) { + for (auto& row : *outerVector) { + row += visitor.numValuesBias(); + } + } bulkScan( folly::Range(rows, outerVector->size()), outerVector->data(), @@ -157,6 +162,11 @@ class RleBpDataDecoder : public bytedance::bolt::parquet::RleBpDecoder { visitor.setAllNull(hasFilter ? 0 : numRows); return; } + if (hasHook && visitor.numValuesBias() > 0) { + for (auto& row : *outerVector) { + row += visitor.numValuesBias(); + } + } bulkScan( *innerVector, outerVector->data(), visitor); skip(tailSkip, 0, nullptr); diff --git a/bolt/dwio/parquet/reader/StringColumnReader.cpp b/bolt/dwio/parquet/reader/StringColumnReader.cpp index 964e2c130..66ed1bbd9 100644 --- a/bolt/dwio/parquet/reader/StringColumnReader.cpp +++ b/bolt/dwio/parquet/reader/StringColumnReader.cpp @@ -137,9 +137,7 @@ void StringColumnReader::read( rows, dwio::common::ExtractToGenericHook(scanSpec_->valueHook())); } - return; - } - if (isDense) { + } else if (isDense) { processFilter( scanSpec_->filter(), rows, dwio::common::ExtractToReader(this)); } else { diff --git a/bolt/dwio/parquet/reader/StringDecoder.h b/bolt/dwio/parquet/reader/StringDecoder.h index a302c08f9..37714e937 100644 --- a/bolt/dwio/parquet/reader/StringDecoder.h +++ b/bolt/dwio/parquet/reader/StringDecoder.h @@ -62,6 +62,7 @@ class StringDecoder { template void readWithVisitor(const uint64_t* FOLLY_NULLABLE nulls, Visitor visitor) { int32_t current = visitor.start(); + int32_t numValues = 0; skip(current, 0, nulls); int32_t toSkip; bool atEnd = false; @@ -76,6 +77,10 @@ class StringDecoder { skip(toSkip, current, nullptr); } if (atEnd) { + if constexpr (Visitor::kHasHook) { + visitor.setNumValues( + Visitor::kHasFilter ? numValues : visitor.numRows()); + } return; } } @@ -85,11 +90,16 @@ class StringDecoder { fixedLength_ > 0 ? readFixedString() : readString(), atEnd); } ++current; + ++numValues; if (toSkip) { skip(toSkip, current, nulls); current += toSkip; } if (atEnd) { + if constexpr (Visitor::kHasHook) { + visitor.setNumValues( + Visitor::kHasFilter ? numValues : visitor.numRows()); + } return; } } diff --git a/bolt/dwio/parquet/reader/StructColumnReader.h b/bolt/dwio/parquet/reader/StructColumnReader.h index 5d73fa3fb..756dcd7fb 100644 --- a/bolt/dwio/parquet/reader/StructColumnReader.h +++ b/bolt/dwio/parquet/reader/StructColumnReader.h @@ -72,6 +72,10 @@ class StructColumnReader : public dwio::common::SelectiveStructColumnReader { void setNullsFromRepDefs(PageReader& pageReader); + bool shouldReadChildrenEagerly() const override { + return fileType_->isDCMap(); + } + dwio::common::SelectiveColumnReader* FOLLY_NULLABLE childForRepDefs() const { return childForRepDefs_; } diff --git a/bolt/dwio/parquet/reader/VariantColumnReader.h b/bolt/dwio/parquet/reader/VariantColumnReader.h index 3e9fd789d..9874c25ba 100644 --- a/bolt/dwio/parquet/reader/VariantColumnReader.h +++ b/bolt/dwio/parquet/reader/VariantColumnReader.h @@ -31,6 +31,10 @@ class VariantColumnReader : public StructColumnReader { memory::MemoryPool& pool); void getValues(const RowSet& rows, VectorPtr* result) override; + + bool shouldReadChildrenEagerly() const override { + return true; + } }; } // namespace bytedance::bolt::parquet diff --git a/bolt/dwio/parquet/tests/ParquetTestBase.h b/bolt/dwio/parquet/tests/ParquetTestBase.h index 0d8fca84b..7e2cfad1a 100644 --- a/bolt/dwio/parquet/tests/ParquetTestBase.h +++ b/bolt/dwio/parquet/tests/ParquetTestBase.h @@ -115,13 +115,14 @@ class ParquetTestBase : public testing::Test, public test::VectorTestBase { const VectorPtr& expected, const VectorPtr& actual, vector_size_t offset) { - ASSERT_GE(expected->size(), actual->size() + offset); - ASSERT_EQ(expected->typeKind(), actual->typeKind()); - for (vector_size_t i = 0; i < actual->size(); i++) { - ASSERT_TRUE(expected->equalValueAt(actual.get(), i + offset, i)) + const auto& loadedActual = BaseVector::loadedVectorShared(actual); + ASSERT_GE(expected->size(), loadedActual->size() + offset); + ASSERT_EQ(expected->typeKind(), loadedActual->typeKind()); + for (vector_size_t i = 0; i < loadedActual->size(); i++) { + ASSERT_TRUE(expected->equalValueAt(loadedActual.get(), i + offset, i)) << "at " << (i + offset) << ": expected " << expected->toString(i + offset) << ", but got " - << actual->toString(i); + << loadedActual->toString(i); } } diff --git a/bolt/dwio/parquet/tests/reader/ParquetReaderTest.cpp b/bolt/dwio/parquet/tests/reader/ParquetReaderTest.cpp index 6c79a6cad..b1162bc97 100644 --- a/bolt/dwio/parquet/tests/reader/ParquetReaderTest.cpp +++ b/bolt/dwio/parquet/tests/reader/ParquetReaderTest.cpp @@ -33,9 +33,7 @@ #include #include #include -#ifdef SPARK_COMPATIBLE #include "bolt/common/base/tests/GTestUtils.h" -#endif #include "bolt/core/QueryCtx.h" #include "bolt/dwio/parquet/reader/RepeatedColumnReader.h" #include "bolt/dwio/parquet/tests/ParquetTestBase.h" @@ -89,25 +87,92 @@ std::string getVariantFixturePath(const std::string& fileName) { cwd.string(), sourceDir.string()); } + +template +T readHookValue(const void* value) { + return *reinterpret_cast(value); +} + +template <> +std::string readHookValue(const void* value) { + return std::string(*reinterpret_cast(value)); +} + +template +class CollectValueHook : public ValueHook { + public: + explicit CollectValueHook(vector_size_t size) + : values_(size), present_(size, false) {} + + void addValue(vector_size_t row, const void* value) override { + present_[row] = true; + values_[row] = readHookValue(value); + } + + const std::vector& values() const { + return values_; + } + + const std::vector& present() const { + return present_; + } + + private: + std::vector values_; + std::vector present_; +}; + +BaseVector* loadedChildAt(RowVector* row, vector_size_t index) { + return row->childAt(index)->loadedVector(); +} + +template +FlatVector* loadedFlatChildAt(RowVector* row, vector_size_t index) { + return loadedChildAt(row, index)->asFlatVector(); +} } // namespace class ParquetReaderTest : public ParquetTestBase { public: - std::unique_ptr createRowReader( - const std::string& fileName, - const RowTypePtr& rowType) { - const std::string sample(getExampleFilePath(fileName)); + std::shared_ptr writeTempParquet( + const std::vector& data, + bytedance::bolt::parquet::WriterOptions options = {}) { + BOLT_CHECK_GT(data.size(), 0); + + auto tempFile = exec::test::TempFilePath::create(); + auto writeFile = + std::make_unique(tempFile->getPath(), true, false); + auto sink = std::make_unique( + std::move(writeFile), tempFile->getPath()); + options.memoryPool = rootPool_.get(); + + auto writer = std::make_unique( + std::move(sink), options, asRowType(data[0]->type())); + for (const auto& batch : data) { + writer->write(batch); + } + writer->close(); + return tempFile; + } + std::unique_ptr createRowReaderFromPath( + const std::string& path, + const RowTypePtr& rowType) { bytedance::bolt::dwio::common::ReaderOptions readerOptions{leafPool_.get()}; - auto reader = createReader(sample, readerOptions); + auto reader = createReader(path, readerOptions); RowReaderOptions rowReaderOpts; rowReaderOpts.select( std::make_shared( rowType, rowType->names())); rowReaderOpts.setScanSpec(makeScanSpec(rowType)); - auto rowReader = reader->createRowReader(rowReaderOpts); - return rowReader; + return reader->createRowReader(rowReaderOpts); + } + + std::unique_ptr createRowReader( + const std::string& fileName, + const RowTypePtr& rowType) { + return createRowReaderFromPath(getExampleFilePath(fileName), rowType); } void assertReadWithExpected( @@ -145,7 +210,7 @@ TEST_F(ParquetReaderTest, parseDecimal) { auto result = BaseVector::create(schema, 1, leafPool_.get()); rowReader->next(6, result); auto decimals = result->as(); - auto a = decimals->childAt(0)->asFlatVector()->rawValues(); + auto a = loadedFlatChildAt(decimals, 0)->rawValues(); EXPECT_EQ(a[0], 64830000); } @@ -729,6 +794,295 @@ TEST_F(ParquetReaderTest, projectNoColumns) { ASSERT_FALSE(rowReader->next(kBatchSize, result)); } +TEST_F(ParquetReaderTest, producesLazyVectorsForProjectedColumns) { + constexpr vector_size_t kSize = 12; + auto data = makeRowVector( + {"a", "b"}, + { + makeFlatVector(kSize, [](auto row) { return row + 1; }), + makeFlatVector(kSize, [](auto row) { return row + 1; }), + }); + auto file = writeTempParquet({data}); + auto rowType = asRowType(data->type()); + auto rowReader = createRowReaderFromPath(file->getPath(), rowType); + auto result = BaseVector::create(rowType, 0, leafPool_.get()); + + ASSERT_TRUE(rowReader->next(5, result)); + auto row = result->asUnchecked(); + ASSERT_EQ(2, row->childrenSize()); + ASSERT_TRUE(isLazyNotLoaded(*row->childAt(0))); + ASSERT_TRUE(isLazyNotLoaded(*row->childAt(1))); + + auto a = row->childAt(0)->loadedVector()->asFlatVector(); + auto b = row->childAt(1)->loadedVector()->asFlatVector(); + ASSERT_EQ(5, row->size()); + for (auto i = 0; i < row->size(); ++i) { + EXPECT_EQ(i + 1, a->valueAt(i)); + EXPECT_EQ(i + 1, b->valueAt(i)); + } +} + +TEST_F(ParquetReaderTest, producesLazyVectorsForComplexProjectedColumns) { + constexpr vector_size_t kSize = 12; + auto data = makeRowVector( + {"id", "items", "payload"}, + { + makeFlatVector(kSize, [](auto row) { return row; }), + makeArrayVector( + kSize, + [](auto row) { return row % 4; }, + [](auto row) { return row * 10; }), + makeRowVector( + {"name", "score"}, + { + makeFlatVector( + kSize, [](auto row) { return fmt::format("n{}", row); }), + makeFlatVector( + kSize, [](auto row) { return row * 1.5; }), + }), + }); + auto file = writeTempParquet({data}); + auto rowType = asRowType(data->type()); + auto rowReader = createRowReaderFromPath(file->getPath(), rowType); + auto result = BaseVector::create(rowType, 0, leafPool_.get()); + + ASSERT_TRUE(rowReader->next(kSize, result)); + auto row = result->asUnchecked(); + ASSERT_EQ(kSize, row->size()); + ASSERT_TRUE(isLazyNotLoaded(*row->childAt(0))); + ASSERT_TRUE(isLazyNotLoaded(*row->childAt(1))); + ASSERT_TRUE(isLazyNotLoaded(*row->childAt(2))); + + for (auto i = 0; i < row->childrenSize(); ++i) { + row->childAt(i)->loadedVector(); + } + test::assertEqualVectors(data, result); +} + +TEST_F(ParquetReaderTest, loadingStaleLazyVectorFailsAfterReaderAdvances) { + constexpr vector_size_t kSize = 12; + auto data = makeRowVector( + {"a", "b"}, + { + makeFlatVector(kSize, [](auto row) { return row + 1; }), + makeFlatVector(kSize, [](auto row) { return row + 0.5; }), + }); + auto file = writeTempParquet({data}); + auto rowType = asRowType(data->type()); + auto rowReader = createRowReaderFromPath(file->getPath(), rowType); + + auto first = BaseVector::create(rowType, 0, leafPool_.get()); + ASSERT_TRUE(rowReader->next(5, first)); + auto staleLazy = first->asUnchecked()->childAt(0); + ASSERT_TRUE(isLazyNotLoaded(*staleLazy)); + + auto second = BaseVector::create(rowType, 0, leafPool_.get()); + ASSERT_TRUE(rowReader->next(5, second)); + BOLT_ASSERT_THROW( + staleLazy->loadedVector(), + "Loading LazyVector after the enclosing reader has moved"); + + auto secondRow = second->asUnchecked(); + ASSERT_TRUE(isLazyNotLoaded(*secondRow->childAt(0))); + auto loaded = secondRow->childAt(0)->loadedVector()->asFlatVector(); + for (auto i = 0; i < secondRow->size(); ++i) { + EXPECT_EQ(i + 6, loaded->valueAt(i)); + } +} + +TEST_F(ParquetReaderTest, lazyLoadAcrossSmallPages) { + constexpr vector_size_t kSize = 32; + auto data = makeRowVector( + {"id", "name"}, + { + makeFlatVector(kSize, [](auto row) { return row * 11; }), + makeFlatVector( + kSize, [](auto row) { return fmt::format("name-{}", row); }), + }); + bytedance::bolt::parquet::WriterOptions options; + options.dataPageSize = 1; + auto file = writeTempParquet({data}, options); + auto rowType = asRowType(data->type()); + auto rowReader = createRowReaderFromPath(file->getPath(), rowType); + auto result = BaseVector::create(rowType, 0, leafPool_.get()); + + ASSERT_TRUE(rowReader->next(kSize, result)); + auto row = result->asUnchecked(); + ASSERT_EQ(kSize, row->size()); + ASSERT_TRUE(isLazyNotLoaded(*row->childAt(0))); + ASSERT_TRUE(isLazyNotLoaded(*row->childAt(1))); + + DecodedVector id(*row->childAt(0), SelectivityVector(kSize)); + DecodedVector name(*row->childAt(1), SelectivityVector(kSize)); + for (auto i = 0; i < kSize; ++i) { + EXPECT_EQ(i * 11, id.valueAt(i)); + EXPECT_EQ(fmt::format("name-{}", i), name.valueAt(i).str()); + } +} + +TEST_F( + ParquetReaderTest, + valueHookSkipsNullsForSparseLazyRowsAcrossSmallPages) { + constexpr vector_size_t kSize = 32; + auto data = makeRowVector( + {"id"}, + {makeFlatVector( + kSize, + [](auto row) { return row * 11; }, + [](auto row) { return row == 3 || row == 14; })}); + bytedance::bolt::parquet::WriterOptions options; + options.enableDictionary = false; + options.dataPageSize = 1; + auto file = writeTempParquet({data}, options); + auto rowType = asRowType(data->type()); + auto rowReader = createRowReaderFromPath(file->getPath(), rowType); + auto result = BaseVector::create(rowType, 0, leafPool_.get()); + + ASSERT_TRUE(rowReader->next(kSize, result)); + auto id = result->asUnchecked()->childAt(0); + ASSERT_TRUE(isLazyNotLoaded(*id)); + + const std::vector selectedRows = {2, 3, 7, 14, 20}; + CollectValueHook hook(kSize); + id->asUnchecked()->load(RowSet(selectedRows), &hook); + + ASSERT_TRUE(hook.present()[0]); + EXPECT_EQ(22, hook.values()[0]); + EXPECT_FALSE(hook.present()[1]); + ASSERT_TRUE(hook.present()[2]); + EXPECT_EQ(77, hook.values()[2]); + EXPECT_FALSE(hook.present()[3]); + ASSERT_TRUE(hook.present()[4]); + EXPECT_EQ(220, hook.values()[4]); + for (auto row = selectedRows.size(); row < kSize; ++row) { + EXPECT_FALSE(hook.present()[row]) << row; + } +} + +TEST_F(ParquetReaderTest, valueHookLoadDoesNotPoisonNextLazyBatch) { + constexpr vector_size_t kSize = 12; + auto data = makeRowVector( + {"id"}, + {makeFlatVector(kSize, [](auto row) { return row * 10; })}); + auto file = writeTempParquet({data}); + auto rowType = asRowType(data->type()); + auto rowReader = createRowReaderFromPath(file->getPath(), rowType); + + auto first = BaseVector::create(rowType, 0, leafPool_.get()); + ASSERT_TRUE(rowReader->next(6, first)); + auto firstId = first->asUnchecked()->childAt(0); + ASSERT_TRUE(isLazyNotLoaded(*firstId)); + + std::vector firstRows(6); + std::iota(firstRows.begin(), firstRows.end(), 0); + CollectValueHook hook(6); + firstId->asUnchecked()->load(RowSet(firstRows), &hook); + for (auto row = 0; row < firstRows.size(); ++row) { + EXPECT_TRUE(hook.present()[row]) << row; + EXPECT_EQ(row * 10, hook.values()[row]); + } + + auto second = BaseVector::create(rowType, 0, leafPool_.get()); + ASSERT_TRUE(rowReader->next(6, second)); + auto secondId = second->asUnchecked()->childAt(0); + ASSERT_TRUE(isLazyNotLoaded(*secondId)); + auto loaded = secondId->loadedVector()->asFlatVector(); + for (auto row = 0; row < second->size(); ++row) { + EXPECT_EQ((row + 6) * 10, loaded->valueAt(row)); + } +} + +TEST_F(ParquetReaderTest, stringValueHookSparseRowsAcrossSmallPages) { + constexpr vector_size_t kSize = 32; + auto data = + makeRowVector({"name"}, {makeFlatVector(kSize, [](auto row) { + return fmt::format("direct-{}", row); + })}); + bytedance::bolt::parquet::WriterOptions options; + options.enableDictionary = false; + options.dataPageSize = 1; + auto file = writeTempParquet({data}, options); + auto rowType = asRowType(data->type()); + auto rowReader = createRowReaderFromPath(file->getPath(), rowType); + auto result = BaseVector::create(rowType, 0, leafPool_.get()); + + ASSERT_TRUE(rowReader->next(kSize, result)); + auto name = result->asUnchecked()->childAt(0); + ASSERT_TRUE(isLazyNotLoaded(*name)); + + const std::vector selectedRows = {0, 5, 17, 31}; + CollectValueHook hook(kSize); + name->asUnchecked()->load(RowSet(selectedRows), &hook); + + for (auto row = 0; row < selectedRows.size(); ++row) { + EXPECT_TRUE(hook.present()[row]) << row; + EXPECT_EQ(fmt::format("direct-{}", selectedRows[row]), hook.values()[row]); + } + for (auto row = selectedRows.size(); row < kSize; ++row) { + EXPECT_FALSE(hook.present()[row]) << row; + } +} + +TEST_F(ParquetReaderTest, stringDictionaryValueHookAcrossSmallPages) { + constexpr vector_size_t kSize = 32; + auto data = + makeRowVector({"name"}, {makeFlatVector(kSize, [](auto row) { + return fmt::format("dict-{}", row % 5); + })}); + bytedance::bolt::parquet::WriterOptions options; + options.enableDictionary = true; + options.dataPageSize = 1; + auto file = writeTempParquet({data}, options); + auto rowType = asRowType(data->type()); + auto rowReader = createRowReaderFromPath(file->getPath(), rowType); + auto result = BaseVector::create(rowType, 0, leafPool_.get()); + + ASSERT_TRUE(rowReader->next(kSize, result)); + auto name = result->asUnchecked()->childAt(0); + ASSERT_TRUE(isLazyNotLoaded(*name)); + + std::vector rows(kSize); + std::iota(rows.begin(), rows.end(), 0); + CollectValueHook hook(kSize); + name->asUnchecked()->load(RowSet(rows), &hook); + + for (auto i = 0; i < kSize; ++i) { + EXPECT_EQ(fmt::format("dict-{}", i % 5), hook.values()[i]); + } +} + +TEST_F(ParquetReaderTest, stringDictionaryValueHookSparseRowsAcrossSmallPages) { + constexpr vector_size_t kSize = 32; + auto data = + makeRowVector({"name"}, {makeFlatVector(kSize, [](auto row) { + return fmt::format("dict-{}", row % 5); + })}); + bytedance::bolt::parquet::WriterOptions options; + options.enableDictionary = true; + options.dataPageSize = 1; + auto file = writeTempParquet({data}, options); + auto rowType = asRowType(data->type()); + auto rowReader = createRowReaderFromPath(file->getPath(), rowType); + auto result = BaseVector::create(rowType, 0, leafPool_.get()); + + ASSERT_TRUE(rowReader->next(kSize, result)); + auto name = result->asUnchecked()->childAt(0); + ASSERT_TRUE(isLazyNotLoaded(*name)); + + const std::vector selectedRows = {1, 3, 7, 14}; + CollectValueHook hook(kSize); + name->asUnchecked()->load(RowSet(selectedRows), &hook); + + for (auto row = 0; row < selectedRows.size(); ++row) { + EXPECT_TRUE(hook.present()[row]) << row; + EXPECT_EQ( + fmt::format("dict-{}", selectedRows[row] % 5), hook.values()[row]); + } + for (auto row = selectedRows.size(); row < kSize; ++row) { + EXPECT_FALSE(hook.present()[row]) << row; + } +} + // Validates the per-row size estimate produced by // ReaderBase::estimatedRowGroupBytesInMemory(), which is summed over the // ScanSpec-projected top-level Parquet column nodes inside @@ -965,8 +1319,8 @@ TEST_F(ParquetReaderTest, parseIntDecimal) { rowReader->next(6, result); EXPECT_EQ(result->size(), 6ULL); auto decimals = result->as(); - auto a = decimals->childAt(0)->asFlatVector()->rawValues(); - auto b = decimals->childAt(1)->asFlatVector()->rawValues(); + auto a = loadedFlatChildAt(decimals, 0)->rawValues(); + auto b = loadedFlatChildAt(decimals, 1)->rawValues(); for (int i = 0; i < 3; i++) { int index = 2 * i; EXPECT_EQ(a[index], expectValues[i]); @@ -1069,7 +1423,8 @@ TEST_F(ParquetReaderTest, parseRowArrayTest) { ASSERT_TRUE(rowReader->next(1, result)); // data: 10, 9, , null, {9}, 2 elements starting at 0 {{9}, {10}}} - auto structArray = result->as()->childAt(5)->as(); + auto structArray = + loadedChildAt(result->as(), 5)->as(); auto structEle = structArray->elements() ->as() ->childAt(0) @@ -1505,7 +1860,7 @@ TEST_F(ParquetReaderTest, readEncryptedParquet) { EXPECT_EQ(rowsRead, 3); EXPECT_EQ(result->size(), 3); - auto ids = result->as()->childAt(0)->asFlatVector(); + auto ids = loadedFlatChildAt(result->as(), 0); EXPECT_EQ(ids->valueAt(0), 1); } @@ -1614,8 +1969,7 @@ TEST_F(ParquetReaderTest, readBinaryAsStringFromNation) { rowReader->next(1, result); EXPECT_EQ( expected, - result->as()->childAt(1)->asFlatVector()->valueAt( - 0)); + loadedFlatChildAt(result->as(), 1)->valueAt(0)); } TEST_F(ParquetReaderTest, readComplexType) { @@ -1689,8 +2043,7 @@ TEST_F(ParquetReaderTest, readFixedLenBinaryAsStringFromUuid) { rowReader->next(1, result); EXPECT_EQ( expected, - result->as()->childAt(0)->asFlatVector()->valueAt( - 0)); + loadedFlatChildAt(result->as(), 0)->valueAt(0)); } TEST_F(ParquetReaderTest, skip) { @@ -1765,8 +2118,7 @@ TEST_F(ParquetReaderTest, readVarbinaryFromFLBA) { rowReader->next(1, result); EXPECT_EQ( expected, - result->as()->childAt(0)->asFlatVector()->valueAt( - 0)); + loadedFlatChildAt(result->as(), 0)->valueAt(0)); } TEST_F(ParquetReaderTest, arrayOfMapOfIntKeyArrayValue) { @@ -2049,7 +2401,7 @@ TEST_F(ParquetReaderTest, integerToVarcharSchemaMismatchCast) { ASSERT_EQ(numRows, 5); auto rowResult = result->as(); - auto colResult = rowResult->childAt(0)->asFlatVector(); + auto colResult = loadedFlatChildAt(rowResult, 0); ASSERT_NE(colResult, nullptr); EXPECT_EQ(colResult->valueAt(0).str(), "1"); EXPECT_EQ(colResult->valueAt(1).str(), "2"); @@ -2157,13 +2509,13 @@ TEST_F(ParquetReaderTest, readVariantParquet) { EXPECT_EQ(result->size(), 4); auto row = result->as(); - auto ids = row->childAt(0)->asFlatVector(); + auto ids = loadedFlatChildAt(row, 0); EXPECT_EQ(ids->valueAt(0), 1); EXPECT_EQ(ids->valueAt(1), 2); EXPECT_EQ(ids->valueAt(2), 3); EXPECT_EQ(ids->valueAt(3), 4); - auto variants = row->childAt(1)->as(); + auto variants = loadedChildAt(row, 1)->as(); std::vector expectedJson = { "{\"a\":1,\"b\":[true,\"x\"],\"c\":{\"d\":3.14}}", "[1,2,3]", @@ -2262,7 +2614,7 @@ TEST_F(ParquetReaderTest, readVariantParquetScanSpecOrderMismatch) { EXPECT_EQ(rowsRead, 4); auto row = result->as(); - auto variants = row->childAt(1)->as(); + auto variants = loadedChildAt(row, 1)->as(); auto variantGet = [&](const VariantValue& value, const StringView& path) { auto out = makeFlatVector(1); @@ -2301,8 +2653,8 @@ TEST_F(ParquetReaderTest, readVariantParquetPrimitivesSpark) { EXPECT_EQ(rowsRead, 12); auto row = result->as(); - auto ids = row->childAt(0)->asFlatVector(); - auto variants = row->childAt(1)->as(); + auto ids = loadedFlatChildAt(row, 0); + auto variants = loadedChildAt(row, 1)->as(); auto decodeValue = [&](const VariantValue& value) -> std::string { if (value.value.empty()) { @@ -2405,9 +2757,9 @@ TEST_F(ParquetReaderTest, readVariantParquetNestedSpark) { EXPECT_EQ(rowsRead, 2); auto row = result->as(); - auto arr = row->childAt(1)->as(); - auto map = row->childAt(2)->as(); - auto st = row->childAt(3)->as(); + auto arr = loadedChildAt(row, 1)->as(); + auto map = loadedChildAt(row, 2)->as(); + auto st = loadedChildAt(row, 3)->as(); auto variantGet = [&](const VariantValue& value, const StringView& path) { auto out = makeFlatVector(1); @@ -2446,8 +2798,8 @@ TEST_F(ParquetReaderTest, readVariantParquetNestedSpark) { EXPECT_EQ(gotK2, "s"); // STRUCT - auto v1 = st->childAt(0)->as()->valueAt(0); - auto v2 = st->childAt(1)->as()->valueAt(0); + auto v1 = loadedChildAt(st, 0)->as()->valueAt(0); + auto v2 = loadedChildAt(st, 1)->as()->valueAt(0); auto [okA, a] = variantGet(v1, "$.a"); EXPECT_TRUE(okA); EXPECT_EQ(a, "1"); @@ -2473,7 +2825,7 @@ TEST_F(ParquetReaderTest, readVariantParquetRawJsonStruct) { EXPECT_EQ(rowsRead, 4); auto row = result->as(); - auto variants = row->childAt(1)->as(); + auto variants = loadedChildAt(row, 1)->as(); // Raw JSON payloads are expected to have empty metadata. EXPECT_TRUE(variants->valueAt(0).metadata.empty()); @@ -2523,7 +2875,7 @@ TEST_F(ParquetReaderTest, readVariantParquetRawParts) { EXPECT_EQ(result->size(), 4); auto row = result->as(); - auto variants = row->childAt(1)->as(); + auto variants = loadedChildAt(row, 1)->as(); for (int i = 0; i < 4; ++i) { if (variants->isNullAt(i)) { diff --git a/bolt/dwio/parquet/tests/reader/ParquetTableScanTest.cpp b/bolt/dwio/parquet/tests/reader/ParquetTableScanTest.cpp index 6e777d158..2b24c47c6 100644 --- a/bolt/dwio/parquet/tests/reader/ParquetTableScanTest.cpp +++ b/bolt/dwio/parquet/tests/reader/ParquetTableScanTest.cpp @@ -67,6 +67,19 @@ std::pair encodeVariantJson(std::string_view json) { return {std::move(value), dict.serialize()}; } +int64_t loadedToValueHook(const std::shared_ptr& task) { + int64_t sum = 0; + for (const auto& pipelineStats : task->taskStats().pipelineStats) { + for (const auto& operatorStats : pipelineStats.operatorStats) { + auto it = operatorStats.runtimeStats.find("loadedToValueHook"); + if (it != operatorStats.runtimeStats.end()) { + sum += it->second.sum; + } + } + } + return sum; +} + RowVectorPtr makeVariantParquetBatch( memory::MemoryPool* pool, const std::vector& groups, @@ -400,6 +413,19 @@ class ParquetTableScanTest : public HiveConnectorTestBase { writer->close(); } + std::unique_ptr makeCursorWithoutCopy( + const core::PlanNodePtr& plan, + const std::string& filePath) { + CursorParameters params; + params.copyResult = false; + params.serialExecution = true; + params.planNode = plan; + auto cursor = TaskCursor::create(params); + cursor->task()->addSplit("0", exec::Split(makeSplit(filePath))); + cursor->task()->noMoreSplits("0"); + return cursor; + } + void testTimestampRead(const WriterOptions& options) { auto stringToTimestamp = [](std::string_view view) { return util::fromTimestampString(view.data(), view.size(), nullptr); @@ -556,6 +582,151 @@ TEST_F(ParquetTableScanTest, basic) { "SELECT max(b), a FROM tmp WHERE a < 3 GROUP BY a"); } +TEST_F(ParquetTableScanTest, tableScanPreservesLazyVectors) { + constexpr vector_size_t kSize = 20; + auto data = makeRowVector( + {"a", "b"}, + { + makeFlatVector(kSize, [](auto row) { return row + 1; }), + makeFlatVector(kSize, [](auto row) { return row + 1; }), + }); + auto file = TempFilePath::create(); + writeToParquetFile(file->getPath(), {data}, {}); + auto rowType = ROW({"a", "b"}, {BIGINT(), DOUBLE()}); + auto plan = PlanBuilder().tableScan(rowType).planNode(); + + auto cursor = makeCursorWithoutCopy(plan, file->getPath()); + vector_size_t rowOffset = 0; + while (cursor->moveNext()) { + auto result = cursor->current()->asUnchecked(); + ASSERT_GT(result->size(), 0); + ASSERT_EQ(2, result->childrenSize()); + ASSERT_TRUE(isLazyNotLoaded(*result->childAt(0))); + ASSERT_TRUE(isLazyNotLoaded(*result->childAt(1))); + + auto a = result->childAt(0)->loadedVector()->asFlatVector(); + auto b = result->childAt(1)->loadedVector()->asFlatVector(); + for (auto i = 0; i < result->size(); ++i) { + EXPECT_EQ(rowOffset + i + 1, a->valueAt(i)); + EXPECT_EQ(rowOffset + i + 1, b->valueAt(i)); + } + rowOffset += result->size(); + } + ASSERT_EQ(20, rowOffset); + ASSERT_TRUE(waitForTaskCompletion(cursor->task().get())); +} + +TEST_F(ParquetTableScanTest, tableScanPreservesComplexLazyVectors) { + constexpr vector_size_t kSize = 8; + auto data = makeRowVector( + {"id", "items", "payload"}, + { + makeFlatVector(kSize, [](auto row) { return row; }), + makeArrayVector( + kSize, + [](auto row) { return row % 3; }, + [](auto row) { return row * 7; }), + makeRowVector( + {"name", "score"}, + { + makeFlatVector( + kSize, [](auto row) { return fmt::format("n{}", row); }), + makeFlatVector( + kSize, [](auto row) { return row * 2.5; }), + }), + }); + auto file = TempFilePath::create(); + writeToParquetFile(file->getPath(), {data}, {}); + auto rowType = asRowType(data->type()); + auto plan = PlanBuilder().tableScan(rowType).planNode(); + + auto cursor = makeCursorWithoutCopy(plan, file->getPath()); + ASSERT_TRUE(cursor->moveNext()); + auto result = cursor->current(); + ASSERT_EQ(kSize, result->size()); + for (auto i = 0; i < result->childrenSize(); ++i) { + ASSERT_TRUE(isLazyNotLoaded(*result->childAt(i))) << i; + result->childAt(i)->loadedVector(); + } + bytedance::bolt::test::assertEqualVectors(data, result); + ASSERT_FALSE(cursor->moveNext()); + ASSERT_TRUE(waitForTaskCompletion(cursor->task().get())); +} + +TEST_F(ParquetTableScanTest, filterColumnsAreEagerButProjectedColumnsAreLazy) { + constexpr vector_size_t kSize = 10; + auto data = makeRowVector( + {"filter_col", "payload", "items"}, + { + makeFlatVector(kSize, [](auto row) { return row; }), + makeFlatVector(kSize, [](auto row) { return row + 0.25; }), + makeArrayVector( + kSize, + [](auto row) { return row % 2; }, + [](auto row) { return row + 100; }), + }); + auto file = TempFilePath::create(); + writeToParquetFile(file->getPath(), {data}, {}); + auto rowType = asRowType(data->type()); + auto plan = PlanBuilder().tableScan(rowType, {"filter_col >= 3"}).planNode(); + + auto cursor = makeCursorWithoutCopy(plan, file->getPath()); + ASSERT_TRUE(cursor->moveNext()); + auto result = cursor->current(); + ASSERT_EQ(7, result->size()); + ASSERT_FALSE(isLazyNotLoaded(*result->childAt(0))); + ASSERT_TRUE(isLazyNotLoaded(*result->childAt(1))); + ASSERT_TRUE(isLazyNotLoaded(*result->childAt(2))); + + auto filterCol = result->childAt(0)->asFlatVector(); + auto payload = result->childAt(1)->loadedVector()->asFlatVector(); + result->childAt(2)->loadedVector(); + for (auto i = 0; i < result->size(); ++i) { + EXPECT_EQ(i + 3, filterCol->valueAt(i)); + EXPECT_EQ(i + 3.25, payload->valueAt(i)); + } + ASSERT_FALSE(cursor->moveNext()); + ASSERT_TRUE(waitForTaskCompletion(cursor->task().get())); +} + +TEST_F(ParquetTableScanTest, remainingFilterPreservesLazyOutputWrapper) { + constexpr vector_size_t kSize = 12; + auto data = makeRowVector( + {"a", "b", "payload"}, + { + makeFlatVector(kSize, [](auto row) { return row; }), + makeFlatVector(kSize, [](auto row) { return row % 3; }), + makeFlatVector( + kSize, [](auto row) { return fmt::format("payload{}", row); }), + }); + auto file = TempFilePath::create(); + writeToParquetFile(file->getPath(), {data}, {}); + auto rowType = asRowType(data->type()); + auto plan = PlanBuilder().tableScan(rowType, {}, "b = 1").planNode(); + + auto cursor = makeCursorWithoutCopy(plan, file->getPath()); + ASSERT_TRUE(cursor->moveNext()); + auto result = cursor->current(); + ASSERT_EQ(4, result->size()); + ASSERT_TRUE(isLazyNotLoaded(*result->childAt(0))); + ASSERT_FALSE(isLazyNotLoaded(*result->childAt(1))); + ASSERT_TRUE(isLazyNotLoaded(*result->childAt(2))); + + DecodedVector a(*result->childAt(0), SelectivityVector(result->size())); + DecodedVector b(*result->childAt(1), SelectivityVector(result->size())); + DecodedVector payload(*result->childAt(2), SelectivityVector(result->size())); + for (auto i = 0; i < result->size(); ++i) { + const auto sourceRow = 1 + 3 * i; + EXPECT_EQ(sourceRow, a.valueAt(i)); + EXPECT_EQ(1, b.valueAt(i)); + EXPECT_EQ( + fmt::format("payload{}", sourceRow), + payload.valueAt(i).str()); + } + ASSERT_FALSE(cursor->moveNext()); + ASSERT_TRUE(waitForTaskCompletion(cursor->task().get())); +} + TEST_F(ParquetTableScanTest, aggregatePushdownToSmallPages) { const std::vector columnNames = {"a", "b", "c"}; const auto expectedRowVector = makeRowVector( @@ -591,6 +762,12 @@ TEST_F(ParquetTableScanTest, aggregatePushdownToSmallPages) { AssertQueryBuilder(plan) .split(makeSplit(filePath->getPath())) .assertResults(expectedRowVector); + + std::shared_ptr task; + AssertQueryBuilder(plan) + .split(makeSplit(filePath->getPath())) + .copyResults(pool(), task); + EXPECT_GT(loadedToValueHook(task), 0); } TEST_F(ParquetTableScanTest, countStar) { diff --git a/bolt/exec/Driver.cpp b/bolt/exec/Driver.cpp index b5a7328f1..11a3f7bda 100644 --- a/bolt/exec/Driver.cpp +++ b/bolt/exec/Driver.cpp @@ -776,8 +776,8 @@ StopReason Driver::runInternal( << " finished. nextOp: " << nextOp->name(); auto timer = createDeltaCpuWallTimer( [nextOp, this](const CpuWallTiming& timing) { - processLazyTiming(*nextOp, timing); - nextOp->stats().wlock()->finishTiming.add(timing); + auto selfDelta = processLazyTiming(*nextOp, timing); + nextOp->stats().wlock()->finishTiming.add(selfDelta); }); BOLT_TEST_ADJUST( "bytedance::bolt::exec::Driver::runInternal::noMoreInput", diff --git a/bolt/exec/FilterProject.cpp b/bolt/exec/FilterProject.cpp index 6cb8d16ca..a013b5b59 100644 --- a/bolt/exec/FilterProject.cpp +++ b/bolt/exec/FilterProject.cpp @@ -54,6 +54,21 @@ bool checkAddIdentityProjection( return false; } +void loadReusedLazyVectors( + const RowVectorPtr& input, + const std::vector& reusedInputChannels) { + if (!input || reusedInputChannels.empty()) { + return; + } + + auto& children = input->children(); + for (auto inputChannel : reusedInputChannels) { + if (isLazyNotLoaded(*children[inputChannel])) { + children[inputChannel]->loadedVector(); + } + } +} + // Split stats to attrbitute cardinality reduction to the Filter node. std::vector splitStats( const OperatorStats& combinedStats, @@ -142,6 +157,17 @@ void FilterProject::initialize() { } isIdentityProjection_ = true; } + { + std::unordered_map inputChannelCounts; + for (const auto& projection : identityProjections_) { + ++inputChannelCounts[projection.inputChannel]; + } + for (const auto& [inputChannel, count] : inputChannelCounts) { + if (count > 1) { + reusedInputChannels_.push_back(inputChannel); + } + } + } numExprs_ = allExprs.size(); exprs_ = makeExprSetFromFlag(std::move(allExprs), operatorCtx_->execCtx()); @@ -169,7 +195,7 @@ void FilterProject::addInput(RowVectorPtr input) { if (!skipForCompositeInput_ || !RowVector::isComposite(input_)) { for (auto& childVec : input_->children()) { if ((acceptCompositeInput_ && childVec == nullptr) || - childVec->isLazy() || + childVec->isLazy() || isLazyNotLoaded(*childVec) || childVec->encoding() != VectorEncoding::Simple::DICTIONARY) { continue; } @@ -240,6 +266,7 @@ RowVectorPtr FilterProject::getOutput() { if (isCompositeInput) { return fillCompositeOutput(size, results); } else { + loadReusedLazyVectors(input_, reusedInputChannels_); return fillOutput(size, nullptr, results); } } @@ -264,6 +291,7 @@ RowVectorPtr FilterProject::getOutput() { results = project(*rows, evalCtx); } + loadReusedLazyVectors(input_, reusedInputChannels_); return fillOutput( numOut, allRowsSelected ? nullptr : filterEvalCtx_.selectedIndices, diff --git a/bolt/exec/FilterProject.h b/bolt/exec/FilterProject.h index 787d51a0c..9b0f6e881 100644 --- a/bolt/exec/FilterProject.h +++ b/bolt/exec/FilterProject.h @@ -149,5 +149,10 @@ class FilterProject : public Operator { // will load c1 only for rows where f(c0) is true. However, c1 identity // projection needs all rows. std::vector multiplyReferencedFieldIndices_; + + // Input channels that map to more than one output channel via identity + // projections. These lazy vectors must be loaded before fillOutput to avoid + // sharing unloaded lazy vectors across multiple output fields. + std::vector reusedInputChannels_; }; } // namespace bytedance::bolt::exec diff --git a/bolt/exec/GroupingSet.cpp b/bolt/exec/GroupingSet.cpp index 2c6e5bad4..10a7ef17a 100644 --- a/bolt/exec/GroupingSet.cpp +++ b/bolt/exec/GroupingSet.cpp @@ -204,6 +204,7 @@ void GroupingSet::addInput(const RowVectorPtr& input, bool mayPushdown) { remainingInput_ = input; firstRemainingRow_ = numRows; remainingMayPushdown_ = mayPushdown; + mayPushdown = false; break; } } diff --git a/bolt/exec/HashProbe.cpp b/bolt/exec/HashProbe.cpp index 32cfdae16..10d4b8cd7 100644 --- a/bolt/exec/HashProbe.cpp +++ b/bolt/exec/HashProbe.cpp @@ -1269,7 +1269,7 @@ RowVectorPtr HashProbe::getOutput() { } } -void HashProbe::fillFilterInput(vector_size_t size) { +RowVectorPtr HashProbe::fillFilterInput(vector_size_t size) { std::vector filterColumns(filterInputType_->size()); for (auto projection : filterInputProjections_) { if (std::any_of( @@ -1305,11 +1305,12 @@ void HashProbe::fillFilterInput(vector_size_t size) { filterInputType_->children(), filterColumns); - filterInput_ = std::make_shared( + return std::make_shared( pool(), filterInputType_, nullptr, size, std::move(filterColumns)); } void HashProbe::prepareFilterRowsForNullAwareJoin( + const RowVector* filterInput, vector_size_t numRows, bool filterPropagateNulls) { BOLT_CHECK_LE(numRows, kBatchSize); @@ -1318,12 +1319,21 @@ void HashProbe::prepareFilterRowsForNullAwareJoin( BaseVector::create(filterInputType_, kBatchSize, pool()); } + if (FOLLY_UNLIKELY(!filterInputRows_.hasSelections())) { + BOLT_CHECK_NULL(filterInput); + if (filterPropagateNulls) { + nullFilterInputRows_.resizeFill(numRows, false); + } + return; + } + BOLT_CHECK_NOT_NULL(filterInput); + if (filterPropagateNulls) { nullFilterInputRows_.resizeFill(numRows, false); auto* rawNullRows = nullFilterInputRows_.asMutableRange().bits(); for (auto& projection : filterInputProjections_) { filterInputColumnDecodedVector_.decode( - *filterInput_->childAt(projection.outputChannel), filterInputRows_); + *filterInput->childAt(projection.outputChannel), filterInputRows_); if (filterInputColumnDecodedVector_.mayHaveNulls()) { SelectivityVector nullsInActiveRows(numRows); memcpy( @@ -1552,16 +1562,20 @@ int32_t HashProbe::evalFilter(int32_t numRows) { filterInputRows_.updateBounds(); } - fillFilterInput(numRows); - - if (nullAware_) { - prepareFilterRowsForNullAwareJoin(numRows, filterPropagateNulls); - } + if (FOLLY_LIKELY(filterInputRows_.hasSelections())) { + auto filterInput = fillFilterInput(numRows); + if (nullAware_) { + prepareFilterRowsForNullAwareJoin( + filterInput.get(), numRows, filterPropagateNulls); + } - EvalCtx evalCtx(operatorCtx_->execCtx(), filter_.get(), filterInput_.get()); - filter_->eval(0, 1, true, filterInputRows_, evalCtx, filterResult_); + EvalCtx evalCtx(operatorCtx_->execCtx(), filter_.get(), filterInput.get()); + filter_->eval(0, 1, true, filterInputRows_, evalCtx, filterResult_); - decodedFilterResult_.decode(*filterResult_[0], filterInputRows_); + decodedFilterResult_.decode(*filterResult_[0], filterInputRows_); + } else if (nullAware_) { + prepareFilterRowsForNullAwareJoin(nullptr, numRows, filterPropagateNulls); + } int32_t numPassed = 0; if (isLeftJoin(joinType_) || isFullJoin(joinType_)) { diff --git a/bolt/exec/HashProbe.h b/bolt/exec/HashProbe.h index 657cea3c4..4caad3c64 100644 --- a/bolt/exec/HashProbe.h +++ b/bolt/exec/HashProbe.h @@ -153,13 +153,14 @@ class HashProbe : public Operator { } // Populate filter input columns. - void fillFilterInput(vector_size_t size); + RowVectorPtr fillFilterInput(vector_size_t size); // Prepare filter row selectivity for null-aware join. 'numRows' // specifies the number of rows in 'filterInputRows_' to process. If // 'filterPropagateNulls' is true, the probe input row which has null in any // probe filter column can't pass the filter. void prepareFilterRowsForNullAwareJoin( + const RowVector* filterInput, vector_size_t numRows, bool filterPropagateNulls); @@ -346,7 +347,7 @@ class HashProbe : public Operator { // side. Used by right semi project join. bool probeSideHasNullKeys_{false}; - // Rows in 'filterInput_' to apply 'filter_' to. + // Rows in the temporary filter input to apply 'filter_' to. SelectivityVector filterInputRows_; // Join filter. @@ -367,11 +368,6 @@ class HashProbe : public Operator { // Maps from column index in hash table to channel in 'filterInputType_'. std::vector filterTableProjections_; - // Temporary projection from probe and build for evaluating - // 'filter_'. This can always be reused since this does not escape - // this operator. - RowVectorPtr filterInput_; - // The following six fields are used in null-aware anti join filter // processing. diff --git a/bolt/exec/OperatorUtils.cpp b/bolt/exec/OperatorUtils.cpp index dc8fd3af8..cec59b40b 100644 --- a/bolt/exec/OperatorUtils.cpp +++ b/bolt/exec/OperatorUtils.cpp @@ -287,8 +287,6 @@ std::vector wrapChildren( std::vector wrappedChildren(children.size()); std::unordered_map old2newMappings; auto curIndex = mapping->as(); - // if children[i]->valueVector()->containingLazyAndWrapped() is true, it means - // the valueVector is isLazyNotLoaded(). for (int i = 0; i < children.size(); ++i) { auto uniqueIter = uniqueDict.find(children[i]); if (uniqueIter != uniqueDict.end()) { @@ -297,12 +295,13 @@ std::vector wrapChildren( } if (nulls == nullptr && children[i]->encoding() == VectorEncoding::Simple::DICTIONARY && - children[i]->rawNulls() == nullptr && - !children[i]->valueVector()->containingLazyAndWrapped()) { + children[i]->rawNulls() == nullptr) { + auto baseValues = children[i]->valueVector(); + baseValues->clearContainingLazyAndWrapped(); auto newMapping = old2newMappings.find(children[i]->wrapInfo()); if (newMapping != old2newMappings.end()) { - wrappedChildren[i] = wrapChild( - size, newMapping->second, children[i]->valueVector(), nullptr); + wrappedChildren[i] = + wrapChild(size, newMapping->second, baseValues, nullptr); } else { // generate new mapping and wrap child.value auto newBuffer = @@ -312,8 +311,7 @@ std::vector wrapChildren( for (auto j = 0; j < size; ++j) { newIndex[j] = childIndex[curIndex[j]]; } - wrappedChildren[i] = - wrapChild(size, newBuffer, children[i]->valueVector(), nullptr); + wrappedChildren[i] = wrapChild(size, newBuffer, baseValues, nullptr); old2newMappings[children[i]->wrapInfo()] = newBuffer; } } else { diff --git a/bolt/exec/OrderBy.cpp b/bolt/exec/OrderBy.cpp index 70c2cd54b..1983248a2 100644 --- a/bolt/exec/OrderBy.cpp +++ b/bolt/exec/OrderBy.cpp @@ -104,6 +104,8 @@ OrderBy::OrderBy( } void OrderBy::addInput(RowVectorPtr input) { + ReclaimableSectionGuard guard(this); + input->loadedVector(); sortBuffer_->addInput(input); } diff --git a/bolt/exec/tests/AggregationTest.cpp b/bolt/exec/tests/AggregationTest.cpp index 4be6cf75c..ed648e7fe 100644 --- a/bolt/exec/tests/AggregationTest.cpp +++ b/bolt/exec/tests/AggregationTest.cpp @@ -49,6 +49,7 @@ #include "bolt/exec/tests/utils/TempDirectoryPath.h" #include "bolt/exec/tests/utils/WithGPUParamInterface.h" #include "bolt/serializers/ArrowSerializer.h" +#include "bolt/vector/LazyVector.h" #include "folly/experimental/EventCount.h" namespace bytedance::bolt::exec::test { @@ -56,6 +57,45 @@ using bytedance::bolt::test::BatchMaker; using core::QueryConfig; using namespace common::testutil; +namespace { +template +class HookOnlyLazyLoader : public VectorLoader { + public: + HookOnlyLazyLoader( + memory::MemoryPool* pool, + std::function valueAt, + vector_size_t size) + : pool_(pool), valueAt_(std::move(valueAt)), size_(size) {} + + private: + void loadInternal( + RowSet rows, + ValueHook* hook, + vector_size_t resultSize, + VectorPtr* result) override { + if (hook != nullptr) { + for (auto row : rows) { + auto value = valueAt_(row); + hook->addValue(row, &value); + } + return; + } + + BaseVector::ensureWritable( + SelectivityVector::empty(), CppToType::create(), pool_, *result); + auto* flat = (*result)->template asFlatVector(); + flat->resize(std::max(resultSize, size_)); + for (auto row : rows) { + flat->set(row, valueAt_(row)); + } + } + + memory::MemoryPool* const pool_; + const std::function valueAt_; + const vector_size_t size_; +}; +} // namespace + /// No-op implementation of Aggregate. Provides public access to following /// base class methods: setNull, clearNull and isNull. class AggregateFunc : public Aggregate { @@ -2938,6 +2978,55 @@ TEST_P(AggregationTest, preGroupedAggregationWithSpilling) { OperatorTestBase::deleteTaskAndCheckSpillDirectory(task); } +TEST_P(AggregationTest, preGroupedSplitWithHookOnlyLazyChild) { + if (GetParam().useGPU) { + GTEST_SKIP() << "GPU Aggregation is not used by this lazy hook test\n"; + } + + constexpr vector_size_t kSize = 1'024; + auto k1 = makeFlatVector(kSize, [](auto row) { return row / 256; }); + auto k2 = makeFlatVector(kSize, [](auto row) { return row % 8; }); + auto valueLazy = std::make_shared( + pool(), + BIGINT(), + kSize, + std::make_unique>( + pool(), + [](vector_size_t row) { return static_cast(row); }, + kSize)); + auto input = makeRowVector({"k1", "k2", "v"}, {k1, k2, valueLazy}); + + auto plan = PlanBuilder() + .values({input}) + .aggregation( + {"k1", "k2"}, + {"k1"}, + {"min(v)", "bitwise_or_agg(v)"}, + {}, + core::AggregationNode::Step::kSingle, + false) + .planNode(); + + std::vector expectedK1(32); + std::vector expectedK2(32); + std::vector expectedMin(32); + std::vector expectedOr(32); + for (vector_size_t group = 0; group < 32; ++group) { + expectedK1[group] = group / 8; + expectedK2[group] = group % 8; + expectedMin[group] = expectedK1[group] * 256 + expectedK2[group]; + expectedOr[group] = (expectedK1[group] << 8) | 248 | expectedK2[group]; + } + auto expected = makeRowVector( + {"k1", "k2", "m", "o"}, + {makeFlatVector(expectedK1), + makeFlatVector(expectedK2), + makeFlatVector(expectedMin), + makeFlatVector(expectedOr)}); + + AssertQueryBuilder(plan).assertResults(expected); +} + TEST_P(AggregationTest, adaptiveOutputBatchRows) { if (GetParam().useGPU) { GTEST_SKIP() << "GPU Aggregation does not support output batch rows"; diff --git a/bolt/exec/tests/FilterProjectTest.cpp b/bolt/exec/tests/FilterProjectTest.cpp index a217cd0f5..8a4128a49 100644 --- a/bolt/exec/tests/FilterProjectTest.cpp +++ b/bolt/exec/tests/FilterProjectTest.cpp @@ -339,6 +339,29 @@ TEST_F(FilterProjectTest, projectAndIdentityOverLazy) { assertQuery(plan, "SELECT c0 < 10 AND c1 < 10, c1 FROM tmp"); } +TEST_F(FilterProjectTest, duplicateIdentityProjectionOverLazy) { + vector_size_t size = 100; + auto valueAt = [](auto row) -> int32_t { return row; }; + auto lazyVectors = makeRowVector({ + makeFlatVector(size, valueAt), + vectorMaker_.lazyFlatVector(size, valueAt), + }); + + auto vectors = makeRowVector({ + makeFlatVector(size, valueAt), + makeFlatVector(size, valueAt), + }); + + createDuckDbTable({vectors}); + + auto plan = PlanBuilder() + .values({lazyVectors}) + .filter("c0 % 2 = 0") + .project({"c1", "c1"}) + .planNode(); + assertQuery(plan, "SELECT c1, c1 FROM tmp WHERE c0 % 2 = 0"); +} + // Verify that nulls on nested parent are propagated to child without copying // the child. Note that null on top level columns are handled separately in // Expr::evalWithNulls; this happens only once per expression tree so we are not diff --git a/bolt/exec/tests/OperatorUtilsTest.cpp b/bolt/exec/tests/OperatorUtilsTest.cpp index 2ecd8e1fa..8ad4921ae 100644 --- a/bolt/exec/tests/OperatorUtilsTest.cpp +++ b/bolt/exec/tests/OperatorUtilsTest.cpp @@ -503,6 +503,47 @@ TEST_F(OperatorUtilsTest, projectDuplicateChildren) { } } +TEST_F(OperatorUtilsTest, projectDictionaryOverLazyCombinesWrappers) { + auto flatVector = + makeFlatVector(5, [](vector_size_t row) { return 10 + row; }); + const auto size = flatVector->size(); + auto lazyVector = std::make_shared( + pool(), + BIGINT(), + size, + std::make_unique( + [&](RowSet /*rows*/) { return flatVector; })); + + auto innerMapping = makeIndices(size, [](vector_size_t row) { + return static_cast(4 - row); + }); + auto dictionaryOverLazy = + BaseVector::wrapInDictionary(nullptr, innerMapping, size, lazyVector); + auto rowVector = makeRowVector({dictionaryOverLazy}); + + auto outerMapping = makeIndices(size, [](vector_size_t row) { + return static_cast((row + 1) % 5); + }); + std::vector projectedChildren(1); + projectChildren( + projectedChildren, + rowVector, + std::vector{{0, 0}}, + size, + outerMapping); + + ASSERT_EQ( + projectedChildren[0]->encoding(), VectorEncoding::Simple::DICTIONARY); + ASSERT_EQ(projectedChildren[0]->valueVector().get(), lazyVector.get()); + + auto expected = makeFlatVector(size, [&](vector_size_t row) { + auto outer = outerMapping->as()[row]; + auto inner = innerMapping->as()[outer]; + return flatVector->valueAt(inner); + }); + bytedance::bolt::test::assertEqualVectors(expected, projectedChildren[0]); +} + TEST_F(OperatorUtilsTest, reclaimableSectionGuard) { RowTypePtr rowType = ROW({"c0"}, {INTEGER()}); diff --git a/bolt/exec/tests/OrderByTest.cpp b/bolt/exec/tests/OrderByTest.cpp index 5facb7acf..993ace462 100644 --- a/bolt/exec/tests/OrderByTest.cpp +++ b/bolt/exec/tests/OrderByTest.cpp @@ -34,6 +34,7 @@ #include #include #include +#include #include #include #include @@ -86,6 +87,38 @@ void abortPool(memory::MemoryPool* pool) { pool->abort(std::current_exception()); } } + +class RecordingLazyLoader : public VectorLoader { + public: + RecordingLazyLoader( + VectorPtr vector, + std::atomic_bool& loaded, + std::atomic_bool& loadedAfterMarker, + std::atomic_bool& marker) + : vector_(std::move(vector)), + loaded_(loaded), + loadedAfterMarker_(loadedAfterMarker), + marker_(marker) {} + + private: + void loadInternal( + RowSet rows, + ValueHook* hook, + vector_size_t resultSize, + VectorPtr* result) override { + BOLT_CHECK(!hook, "RecordingLazyLoader doesn't support ValueHook"); + loaded_ = true; + loadedAfterMarker_ = loadedAfterMarker_ || marker_.load(); + + BOLT_CHECK_EQ(rows.size(), vector_->size()); + *result = BaseVector::copy(*vector_); + } + + const VectorPtr vector_; + std::atomic_bool& loaded_; + std::atomic_bool& loadedAfterMarker_; + std::atomic_bool& marker_; +}; } // namespace class OrderByTest : public OperatorTestBase, public WithGPUParamInterface<> { @@ -2136,6 +2169,57 @@ DEBUG_ONLY_TEST_P(OrderByTest, reclaimFromEmptyOrderBy) { ASSERT_EQ(stats[0].operatorStats[1].spilledPartitions, 0); } +DEBUG_ONLY_TEST_P(OrderByTest, orderByWithLazyInput) { + if (GetParam().useGPU) { + GTEST_SKIP() << "GPU OrderBy is not used by this lazy input test\n"; + } + + auto nonLazyVector = createVectors(1, rowType_, fuzzerOpts_)[0]; + std::atomic_bool sortBufferAddInputEntered{false}; + std::atomic_bool lazyLoaded{false}; + std::atomic_bool lazyLoadedInSortBufferAddInput{false}; + + std::vector lazyChildren; + for (const auto& child : nonLazyVector->children()) { + lazyChildren.push_back(std::make_shared( + pool(), + child->type(), + child->size(), + std::make_unique( + child, + lazyLoaded, + lazyLoadedInSortBufferAddInput, + sortBufferAddInputEntered))); + } + auto lazyInput = std::make_shared( + pool(), + rowType_, + nullptr, + nonLazyVector->size(), + std::move(lazyChildren)); + + createDuckDbTable({nonLazyVector}); + + SCOPED_TESTVALUE_SET( + "bytedance::bolt::exec::SortBuffer::addInput", + std::function( + ([&](void* /*unused*/) { sortBufferAddInputEntered = true; }))); + + const auto spillDirectory = exec::test::TempDirectoryPath::create(); + AssertQueryBuilder(duckDbQueryRunner_) + .spillDirectory(spillDirectory->path) + .config(core::QueryConfig::kSpillEnabled, true) + .config(core::QueryConfig::kOrderBySpillEnabled, true) + .plan(PlanBuilder() + .values({lazyInput}) + .orderBy({"c0 ASC NULLS LAST"}, false) + .planNode()) + .assertResults("SELECT * FROM tmp ORDER BY c0 ASC NULLS LAST"); + + ASSERT_TRUE(lazyLoaded); + ASSERT_FALSE(lazyLoadedInSortBufferAddInput); +} + INSTANTIATE_GPU_TEST_SUITE_P(OrderByTestOnCPUOrGPU, OrderByTest); } // namespace bytedance::bolt::exec::test diff --git a/bolt/exec/tests/utils/QueryAssertions.cpp b/bolt/exec/tests/utils/QueryAssertions.cpp index 2e5b7fcb9..b3dd13a4d 100644 --- a/bolt/exec/tests/utils/QueryAssertions.cpp +++ b/bolt/exec/tests/utils/QueryAssertions.cpp @@ -899,13 +899,19 @@ std::vector materialize(const RowVectorPtr& vector) { rows.reserve(size); auto rowType = vector->type()->as(); + std::vector loadedChildren; + loadedChildren.reserve(rowType.size()); + for (size_t i = 0; i < rowType.size(); ++i) { + loadedChildren.push_back( + BaseVector::loadedVectorShared(vector->childAt(i))); + } for (size_t i = 0; i < size; ++i) { auto numColumns = rowType.size(); MaterializedRow row; row.reserve(numColumns); for (size_t j = 0; j < numColumns; ++j) { - row.push_back(variantAt(vector->childAt(j), i)); + row.push_back(variantAt(loadedChildren[j], i)); } rows.push_back(row); } @@ -1418,7 +1424,9 @@ std::pair, std::vector> readCursor( addSplits(task); while (cursor->moveNext()) { - result.push_back(cursor->current()); + auto vector = cursor->current(); + vector->loadedVector(); + result.push_back(std::move(vector)); addSplits(task); } diff --git a/bolt/vector/BaseVector.cpp b/bolt/vector/BaseVector.cpp index de7ab7244..06feb4f5b 100644 --- a/bolt/vector/BaseVector.cpp +++ b/bolt/vector/BaseVector.cpp @@ -881,9 +881,8 @@ void BaseVector::flattenVector(VectorPtr& vector) { return; } case VectorEncoding::Simple::LAZY: { - auto& loadedVector = - vector->asUnchecked()->loadedVectorShared(); - BaseVector::flattenVector(loadedVector); + vector = vector->asUnchecked()->loadedVectorShared(); + BaseVector::flattenVector(vector); return; } default: diff --git a/bolt/vector/BaseVector.h b/bolt/vector/BaseVector.h index c684e80c5..a8444ef3a 100644 --- a/bolt/vector/BaseVector.h +++ b/bolt/vector/BaseVector.h @@ -36,6 +36,7 @@ #include #include #include +#include #include #include #include @@ -978,7 +979,7 @@ class BaseVector { /// unloaded lazy vector should not be wrapped by two separate top level /// vectors. This would ensure we avoid it being loaded for two separate set /// of rows. - bool containsLazyAndIsWrapped_{false}; + std::atomic_bool containsLazyAndIsWrapped_{false}; // Whether we should use Expr::evalWithMemo to cache the result of evaluation // on dictionary values (this vector). Set to false when the dictionary diff --git a/bolt/vector/ConstantVector.h b/bolt/vector/ConstantVector.h index 51a033192..f9b6313b6 100644 --- a/bolt/vector/ConstantVector.h +++ b/bolt/vector/ConstantVector.h @@ -143,7 +143,11 @@ class ConstantVector final : public SimpleVector { setInternalState(); } - virtual ~ConstantVector() override = default; + virtual ~ConstantVector() override { + if (valueVector_) { + valueVector_->clearContainingLazyAndWrapped(); + } + } bool isNullAt(vector_size_t /*idx*/) const override { BOLT_DCHECK(initialized_); diff --git a/bolt/vector/DictionaryVector.h b/bolt/vector/DictionaryVector.h index b670206dd..2964f77ad 100644 --- a/bolt/vector/DictionaryVector.h +++ b/bolt/vector/DictionaryVector.h @@ -78,7 +78,11 @@ class DictionaryVector : public SimpleVector { std::optional representedBytes = std::nullopt, std::optional storageByteCount = std::nullopt); - virtual ~DictionaryVector() override = default; + virtual ~DictionaryVector() override { + if (dictionaryValues_) { + dictionaryValues_->clearContainingLazyAndWrapped(); + } + } bool mayHaveNulls() const override { BOLT_DCHECK(initialized_); @@ -225,7 +229,8 @@ class DictionaryVector : public SimpleVector { } void setDictionaryValues(VectorPtr dictionaryValues) { - dictionaryValues_ = dictionaryValues; + dictionaryValues_->clearContainingLazyAndWrapped(); + dictionaryValues_ = std::move(dictionaryValues); initialized_ = false; setInternalState(); } diff --git a/bolt/vector/LazyVector.h b/bolt/vector/LazyVector.h index 6a49ce066..ef0f0ac88 100644 --- a/bolt/vector/LazyVector.h +++ b/bolt/vector/LazyVector.h @@ -158,6 +158,7 @@ class LazyVector : public BaseVector { loader_ = std::move(loader); allLoaded_ = false; containsLazyAndIsWrapped_ = false; + resetNulls(); } inline bool isLoaded() const { diff --git a/bolt/vector/SimpleVector.h b/bolt/vector/SimpleVector.h index 204db2911..9b539f934 100644 --- a/bolt/vector/SimpleVector.h +++ b/bolt/vector/SimpleVector.h @@ -158,6 +158,9 @@ class SimpleVector : public BaseVector { vector_size_t index, vector_size_t otherIndex, CompareFlags flags) const override { + // 'this' cannot be a LazyVector, but it may be a dictionary or constant + // wrapping a lazy vector. Load it before accessing values. + loadedVector(); other = other->loadedVector(); DCHECK(dynamic_cast*>(other) != nullptr) << "Attempting to compare vectors not of the same type"; diff --git a/bolt/vector/arrow/Bridge.cpp b/bolt/vector/arrow/Bridge.cpp index 2e17f5f18..4be69e429 100644 --- a/bolt/vector/arrow/Bridge.cpp +++ b/bolt/vector/arrow/Bridge.cpp @@ -2170,6 +2170,11 @@ void exportToArrow( ArrowArray& arrowArray, memory::MemoryPool* pool, const ArrowOptions& options) { + if (vector->encoding() == VectorEncoding::Simple::LAZY) { + exportToArrow( + BaseVector::loadedVectorShared(vector), arrowArray, pool, options); + return; + } if (vector->encoding() == VectorEncoding::Simple::CONSTANT && options.flattenConstant && vector->valueVector() != nullptr && !vector->wrappedVector()->isFlatEncoding()) { @@ -2189,6 +2194,16 @@ void exportToArrow( const ArrowOptions& options, const std::vector& fieldNames, memory::MemoryPool* pool) { + if (vec->encoding() == VectorEncoding::Simple::LAZY) { + exportToArrow( + BaseVector::loadedVectorShared(vec), + arrowSchema, + options, + fieldNames, + pool); + return; + } + auto& type = vec->type(); arrowSchema.name = nullptr; diff --git a/bolt/vector/arrow/tests/ArrowBridgeArrayTest.cpp b/bolt/vector/arrow/tests/ArrowBridgeArrayTest.cpp index 3762a7305..21c527a9b 100644 --- a/bolt/vector/arrow/tests/ArrowBridgeArrayTest.cpp +++ b/bolt/vector/arrow/tests/ArrowBridgeArrayTest.cpp @@ -40,6 +40,7 @@ #include "bolt/vector/arrow/Bridge.h" #include "bolt/vector/tests/utils/VectorMaker.h" #include "bolt/vector/tests/utils/VectorTestBase.h" + namespace bytedance::bolt::test { namespace { @@ -621,6 +622,51 @@ TEST_F(ArrowBridgeArrayExportTest, rowVector) { EXPECT_EQ(nullptr, arrowArray.private_data); } +TEST_F(ArrowBridgeArrayExportTest, rowVectorWithLazyStringChild) { + constexpr vector_size_t kSize = 4; + auto lazyString = std::make_shared( + pool_.get(), + VARCHAR(), + kSize, + std::make_unique([&](RowSet /*rows*/) { + return BaseVector::wrapInDictionary( + nullptr, + makeIndices( + kSize, + [](vector_size_t row) { + static constexpr vector_size_t kIndices[] = {0, 1, 0, 2}; + return kIndices[row]; + }, + pool_.get()), + kSize, + vectorMaker_.flatVector( + {"alpha", "long string value", "gamma"})); + })); + auto vector = vectorMaker_.rowVector({lazyString}); + + ArrowSchema arrowSchema; + ArrowArray arrowArray; + bolt::exportToArrow(vector, arrowSchema, options_); + bolt::exportToArrow(vector, arrowArray, pool_.get(), options_); + + ASSERT_EQ(1, arrowSchema.n_children); + ASSERT_EQ(1, arrowArray.n_children); + ASSERT_NE(nullptr, arrowSchema.children[0]); + ASSERT_NE(nullptr, arrowArray.children[0]); + EXPECT_STREQ("i", arrowSchema.children[0]->format); + ASSERT_NE(nullptr, arrowSchema.children[0]->dictionary); + ASSERT_NE(nullptr, arrowArray.children[0]->dictionary); + EXPECT_STREQ("u", arrowSchema.children[0]->dictionary->format); + EXPECT_EQ(2, arrowArray.children[0]->n_buffers); + EXPECT_EQ(3, arrowArray.children[0]->dictionary->n_buffers); + + auto imported = bolt::importFromArrowAsOwner( + arrowSchema, arrowArray, options_, pool_.get()); + test::assertEqualVectors(BaseVector::loadedVectorShared(vector), imported); + EXPECT_EQ(nullptr, arrowSchema.release); + EXPECT_EQ(nullptr, arrowArray.release); +} + // Test a rowVector containing null entries (in the parent RowVector). TEST_F(ArrowBridgeArrayExportTest, rowVectorNullable) { std::vector> col1 = {0, 1, 2, 3, 4, 5, 6, 7, 8, 9}; diff --git a/bolt/vector/tests/LazyVectorTest.cpp b/bolt/vector/tests/LazyVectorTest.cpp index 3e07bd30d..e17f089b1 100644 --- a/bolt/vector/tests/LazyVectorTest.cpp +++ b/bolt/vector/tests/LazyVectorTest.cpp @@ -83,6 +83,68 @@ TEST_F(LazyVectorTest, lazyInDictionary) { assertCopyableVector(wrapped); } +TEST_F(LazyVectorTest, compareLazies) { + constexpr vector_size_t kInnerSize = 100; + constexpr vector_size_t kOuterSize = 100; + + auto makeLazy = [&]() { + return std::make_shared( + pool_.get(), + INTEGER(), + kInnerSize, + std::make_unique([&](RowSet /*rows*/) { + return makeFlatVector( + kInnerSize, [](vector_size_t row) { return row; }); + })); + }; + + auto makeDictionary = [&]() { + return BaseVector::wrapInDictionary( + nullptr, + makeIndices(kOuterSize, [](vector_size_t row) { return row; }), + kOuterSize, + makeLazy()); + }; + + auto expected = makeFlatVector( + kInnerSize, [](vector_size_t row) { return row; }); + + auto lazy1 = makeLazy(); + auto lazy2 = makeLazy(); + for (vector_size_t i = 0; i < kInnerSize; ++i) { + EXPECT_EQ(lazy1->compare(expected.get(), i, i, CompareFlags()), 0); + EXPECT_EQ(expected->compare(lazy2.get(), i, i, CompareFlags()), 0); + } + + auto dictionaryLazy1 = makeDictionary(); + auto dictionaryLazy2 = makeDictionary(); + for (vector_size_t i = 0; i < kOuterSize; ++i) { + EXPECT_EQ( + dictionaryLazy1->compare(expected.get(), i, i, CompareFlags()), 0); + EXPECT_EQ( + expected->compare(dictionaryLazy2.get(), i, i, CompareFlags()), 0); + } +} + +TEST_F(LazyVectorTest, resetClearsNulls) { + constexpr vector_size_t kSize = 10; + auto loader = [&](RowSet rows) { + return makeFlatVector( + rows.back() + 1, [](vector_size_t row) { return row; }); + }; + LazyVector lazy( + pool_.get(), + INTEGER(), + kSize, + std::make_unique(loader)); + + lazy.setNull(0, true); + ASSERT_NE(lazy.nulls(), nullptr); + + lazy.reset(std::make_unique(loader), kSize); + ASSERT_EQ(lazy.nulls(), nullptr); +} + TEST_F(LazyVectorTest, rowVectorWithLazyChild) { constexpr vector_size_t size = 1000; auto columnType = ROW({"a", "b"}, {INTEGER(), INTEGER()}); diff --git a/bolt/vector/tests/VectorTest.cpp b/bolt/vector/tests/VectorTest.cpp index 40130903c..cda81cc5e 100644 --- a/bolt/vector/tests/VectorTest.cpp +++ b/bolt/vector/tests/VectorTest.cpp @@ -2301,6 +2301,20 @@ TEST_F(VectorTest, nestedLazy) { EXPECT_NO_THROW(BaseVector::wrapInDictionary( nullptr, makeIndices(size, indexAt), size, dict)); + // Verify that if the original dictionary layer is destroyed without loading + // the underlying vector then the lazy vector can be wrapped in a new encoding + // layer. + dict.reset(); + EXPECT_NO_THROW( + dict = BaseVector::wrapInDictionary( + nullptr, makeIndices(size, indexAt), size, lazy)); + + auto replacement = makeLazy(); + auto* rawDict = dict->as>(); + EXPECT_NO_THROW(rawDict->setDictionaryValues(replacement)); + EXPECT_NO_THROW(BaseVector::wrapInDictionary( + nullptr, makeIndices(size, indexAt), size, lazy)); + // Limitation: Current checks cannot prevent existing references of the lazy // vector to load rows. For example, the following would succeed even though // it was wrapped under 2 layer of dictionary above: @@ -2793,24 +2807,44 @@ TEST_F(VectorTest, flattenVector) { VectorPtr lazy = std::make_shared( pool(), INTEGER(), flat->size(), std::make_unique(flat)); - test(lazy, true); + test(lazy, false); - // Constant + // Constant. VectorPtr constant = BaseVector::wrapInConstant(100, 1, flat); test(constant, false); EXPECT_TRUE(constant->isFlatEncoding()); - // Dictionary + // Dictionary. VectorPtr dictionary = BaseVector::wrapInDictionary( nullptr, makeIndices(100, [](auto row) { return row % 2; }), 100, flat); test(dictionary, false); EXPECT_TRUE(dictionary->isFlatEncoding()); + // Lazy dictionary. VectorPtr lazyDictionary = wrapInLazyDictionary(makeFlatVector({1, 2, 3})); - test(lazyDictionary, true); - EXPECT_TRUE(lazyDictionary->isLazy()); - EXPECT_TRUE(lazyDictionary->loadedVector()->isFlatEncoding()); + test(lazyDictionary, false); + EXPECT_TRUE(lazyDictionary->isFlatEncoding()); + + // Lazy dictionary row, where children are also lazy dictionaries. + VectorPtr rowWithLazyDictionaryChildren = wrapInLazyDictionary(makeRowVector({ + wrapInLazyDictionary(makeFlatVector({1, 2, 3})), + wrapInLazyDictionary(array), + wrapInLazyDictionary(map), + })); + test(rowWithLazyDictionaryChildren, false); + EXPECT_EQ( + rowWithLazyDictionaryChildren->encoding(), VectorEncoding::Simple::ROW); + auto* rowWithLazyDictionaryChildrenRowVector = + rowWithLazyDictionaryChildren->as(); + EXPECT_TRUE( + rowWithLazyDictionaryChildrenRowVector->childAt(0)->isFlatEncoding()); + EXPECT_EQ( + rowWithLazyDictionaryChildrenRowVector->childAt(1)->encoding(), + VectorEncoding::Simple::ARRAY); + EXPECT_EQ( + rowWithLazyDictionaryChildrenRowVector->childAt(2)->encoding(), + VectorEncoding::Simple::MAP); // Array with constant elements. auto* arrayVector = array->as();