Skip to content

Commit 4808bf7

Browse files
xin-zhang2xin-zhang2
committed
feat(PartitionedOutput): Add StringView support in PrestoIterativePartitioningSerializer
1 parent e1cc176 commit 4808bf7

4 files changed

Lines changed: 381 additions & 19 deletions

‎velox/serializers/PrestoIterativePartitioningSerializer.cpp‎

Lines changed: 204 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -127,6 +127,72 @@ simpleColumnBytes(const TypePtr& colType, int64_t numRows, int64_t numNulls) {
127127
(numRows - numNulls) * fixedTypeWidth(colType->kind()); // values
128128
}
129129

130+
int64_t variableWidthColumnBytes(
131+
int64_t numRows,
132+
int64_t numNulls,
133+
int64_t valueBytes) {
134+
return 4 + static_cast<int64_t>(kVariableWidth.size()) + // header
135+
4 + // rowCount
136+
4 * numRows + // offsets
137+
1 + // nullFlag
138+
(numNulls > 0 ? bits::nbytes(numRows) : 0) + // null bitmap
139+
4 + // valueBytes
140+
valueBytes; // values
141+
}
142+
143+
std::vector<int64_t> computeVariableWidthDataSizes(
144+
const std::vector<PartitionedVectorPtr>& partitionedVectors,
145+
uint32_t numPartitions) {
146+
std::vector<int64_t> dataSizes(numPartitions, 0);
147+
for (const auto& partitionedVector : partitionedVectors) {
148+
auto encoding = partitionedVector->baseVector()->encoding();
149+
switch (encoding) {
150+
case VectorEncoding::Simple::FLAT: {
151+
auto* flatVector =
152+
partitionedVector->as<PartitionedFlatVector<StringView>>();
153+
VELOX_DCHECK_NOT_NULL(flatVector);
154+
155+
const auto* rawValues =
156+
flatVector->baseVector()->asFlatVector<StringView>()->rawValues();
157+
const auto* rawNulls =
158+
flatVector->baseVector()->asFlatVector<StringView>()->rawNulls();
159+
const auto* rawPartitionOffsets =
160+
partitionedVector->rawPartitionOffsets();
161+
162+
vector_size_t lastOffset = 0;
163+
for (uint32_t p = 0; p < numPartitions; ++p) {
164+
const auto offset = rawPartitionOffsets[p];
165+
if (!rawNulls) {
166+
for (auto i = lastOffset; i < offset; ++i) {
167+
dataSizes[p] += rawValues[i].size();
168+
}
169+
} else {
170+
for (auto i = lastOffset; i < offset; ++i) {
171+
if (!bits::isBitNull(rawNulls, i)) {
172+
dataSizes[p] += rawValues[i].size();
173+
}
174+
}
175+
}
176+
lastOffset = offset;
177+
}
178+
break;
179+
}
180+
case VectorEncoding::Simple::BIASED:
181+
case VectorEncoding::Simple::CONSTANT:
182+
case VectorEncoding::Simple::DICTIONARY:
183+
case VectorEncoding::Simple::SEQUENCE:
184+
VELOX_NYI(
185+
"Unsupported vector encoding for PrestoIterativePartitioningSerializer: ",
186+
encoding);
187+
default:
188+
VELOX_UNSUPPORTED(
189+
"Invalid vector encoding for PrestoIterativePartitioningSerializer: ",
190+
encoding);
191+
}
192+
}
193+
return dataSizes;
194+
}
195+
130196
/// Returns per-partition exact byte counts for one column (all partitions).
131197
/// Recurses into nested ROW columns.
132198
///
@@ -167,9 +233,17 @@ std::vector<int64_t> computeColumnFlushSizes(
167233
}
168234
break;
169235

170-
case TypeKind::TIMESTAMP:
171236
case TypeKind::VARCHAR:
172-
case TypeKind::VARBINARY:
237+
case TypeKind::VARBINARY: {
238+
const auto dataSizes =
239+
computeVariableWidthDataSizes(columnVectors, numPartitions);
240+
for (uint32_t p : nonEmptyPartitions) {
241+
sizes[p] = variableWidthColumnBytes(
242+
rowsPerPartition[p], nullCounts[p], dataSizes[p]);
243+
}
244+
break;
245+
}
246+
case TypeKind::TIMESTAMP:
173247
case TypeKind::ARRAY:
174248
case TypeKind::MAP:
175249
VELOX_NYI(
@@ -478,9 +552,12 @@ void PrestoIterativePartitioningSerializer::flushColumn(
478552
partitionedVectors, colType, nonEmptyPartitions, outputStreams);
479553
break;
480554

481-
case TypeKind::TIMESTAMP:
482555
case TypeKind::VARCHAR:
483556
case TypeKind::VARBINARY:
557+
flushVariableWidthColumn(
558+
partitionedVectors, colType, nonEmptyPartitions, outputStreams);
559+
break;
560+
case TypeKind::TIMESTAMP:
484561
case TypeKind::ROW:
485562
case TypeKind::ARRAY:
486563
case TypeKind::MAP:
@@ -507,6 +584,25 @@ void PrestoIterativePartitioningSerializer::flushSimpleColumn(
507584
}
508585
}
509586

587+
void PrestoIterativePartitioningSerializer::flushVariableWidthColumn(
588+
const std::vector<PartitionedVectorPtr>& partitionedVectors,
589+
const TypePtr& colType,
590+
const std::vector<uint32_t>& nonEmptyPartitions,
591+
const std::vector<IOBufOutputStream*>& outputStreams) const {
592+
flushHeader(typeToEncodingName(colType), nonEmptyPartitions, outputStreams);
593+
flushRowCounts(nonEmptyPartitions, outputStreams);
594+
const auto dataSizes = flushOffsets(partitionedVectors, outputStreams);
595+
flushNulls(partitionedVectors, nonEmptyPartitions, outputStreams);
596+
597+
for (uint32_t p : nonEmptyPartitions) {
598+
writeInt32(outputStreams[p], dataSizes[p]);
599+
}
600+
601+
for (size_t i = 0; i < partitionedVectors.size(); i++) {
602+
flushSingleVariableWidthVector(partitionedVectors[i], outputStreams);
603+
}
604+
}
605+
510606
template <TypeKind kind>
511607
void PrestoIterativePartitioningSerializer::flushSingleFlatVector(
512608
const PartitionedVectorPtr& partitionedVector,
@@ -590,6 +686,60 @@ void PrestoIterativePartitioningSerializer::flushSingleSimpleVector(
590686
}
591687
}
592688

689+
void PrestoIterativePartitioningSerializer::flushSingleVariableWidthFlatVector(
690+
const PartitionedVectorPtr& partitionedVector,
691+
const std::vector<IOBufOutputStream*>& outputStreams) const {
692+
auto* flatVector = partitionedVector->as<PartitionedFlatVector<StringView>>();
693+
VELOX_DCHECK_NOT_NULL(flatVector);
694+
const auto* rawValues =
695+
flatVector->baseVector()->as<FlatVector<StringView>>()->rawValues();
696+
const auto* rawNulls = flatVector->baseVector()->rawNulls();
697+
const auto* partitionOffsets = flatVector->rawPartitionOffsets();
698+
699+
vector_size_t lastOffset = 0;
700+
for (uint32_t p = 0; p < numPartitions_; ++p) {
701+
const auto offset = partitionOffsets[p];
702+
if (outputStreams[p] != nullptr) {
703+
if (!rawNulls) {
704+
for (auto i = lastOffset; i < offset; ++i) {
705+
const auto& value = rawValues[i];
706+
outputStreams[p]->write(value.data(), value.size());
707+
}
708+
} else {
709+
for (auto i = lastOffset; i < offset; ++i) {
710+
if (!bits::isBitNull(rawNulls, i)) {
711+
const auto& value = rawValues[i];
712+
outputStreams[p]->write(value.data(), value.size());
713+
}
714+
}
715+
}
716+
lastOffset = offset;
717+
}
718+
}
719+
}
720+
721+
void PrestoIterativePartitioningSerializer::flushSingleVariableWidthVector(
722+
const PartitionedVectorPtr& partitionedVector,
723+
const std::vector<IOBufOutputStream*>& outputStreams) const {
724+
auto encoding = partitionedVector->baseVector()->encoding();
725+
switch (encoding) {
726+
case VectorEncoding::Simple::FLAT:
727+
flushSingleVariableWidthFlatVector(partitionedVector, outputStreams);
728+
break;
729+
case VectorEncoding::Simple::BIASED:
730+
case VectorEncoding::Simple::CONSTANT:
731+
case VectorEncoding::Simple::DICTIONARY:
732+
case VectorEncoding::Simple::SEQUENCE:
733+
VELOX_NYI(
734+
"Unsupported vector encoding for PrestoIterativePartitioningSerializer: ",
735+
encoding);
736+
default:
737+
VELOX_UNSUPPORTED(
738+
"Invalid vector encoding for PrestoIterativePartitioningSerializer: ",
739+
encoding);
740+
}
741+
}
742+
593743
// ---------------------------------------------------------------------------
594744
// Column building blocks
595745
// ---------------------------------------------------------------------------
@@ -729,4 +879,55 @@ void PrestoIterativePartitioningSerializer::flushSequentialOffsets(
729879
}
730880
}
731881

882+
std::vector<int32_t> PrestoIterativePartitioningSerializer::flushOffsets(
883+
const std::vector<PartitionedVectorPtr>& partitionedVectors,
884+
const std::vector<IOBufOutputStream*>& outputStreams) const {
885+
std::vector<int32_t> offsets(numPartitions_, 0);
886+
for (const auto& partitionedVector : partitionedVectors) {
887+
auto encoding = partitionedVector->baseVector()->encoding();
888+
switch (encoding) {
889+
case VectorEncoding::Simple::FLAT: {
890+
auto* flatVector =
891+
partitionedVector->as<PartitionedFlatVector<StringView>>();
892+
VELOX_DCHECK_NOT_NULL(flatVector);
893+
const auto* rawValues =
894+
flatVector->baseVector()->as<FlatVector<StringView>>()->rawValues();
895+
const auto* rawNulls = flatVector->baseVector()->rawNulls();
896+
const auto* partitionOffsets = flatVector->rawPartitionOffsets();
897+
vector_size_t lastOffset = 0;
898+
for (uint32_t p = 0; p < numPartitions_; ++p) {
899+
const auto offset = partitionOffsets[p];
900+
if (outputStreams[p] != nullptr) {
901+
if (!rawNulls) {
902+
for (auto i = lastOffset; i < offset; ++i) {
903+
offsets[p] += rawValues[i].size();
904+
writeInt32(outputStreams[p], offsets[p]);
905+
}
906+
} else {
907+
for (auto i = lastOffset; i < offset; ++i) {
908+
if (!bits::isBitNull(rawNulls, i)) {
909+
offsets[p] += rawValues[i].size();
910+
}
911+
writeInt32(outputStreams[p], offsets[p]);
912+
}
913+
}
914+
}
915+
lastOffset = offset;
916+
}
917+
} break;
918+
case VectorEncoding::Simple::BIASED:
919+
case VectorEncoding::Simple::CONSTANT:
920+
case VectorEncoding::Simple::DICTIONARY:
921+
case VectorEncoding::Simple::SEQUENCE:
922+
VELOX_NYI(
923+
"Unsupported vector encoding for PrestoIterativePartitioningSerializer: ",
924+
encoding);
925+
default:
926+
VELOX_UNSUPPORTED(
927+
"Invalid vector encoding for PrestoIterativePartitioningSerializer: ",
928+
encoding);
929+
}
930+
}
931+
return offsets;
932+
}
732933
} // namespace facebook::velox::serializer::presto

‎velox/serializers/PrestoIterativePartitioningSerializer.h‎

Lines changed: 18 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -107,15 +107,29 @@ class PrestoIterativePartitioningSerializer {
107107
const std::vector<uint32_t>& nonEmptyPartitions,
108108
const std::vector<IOBufOutputStream*>& outputStreams) const;
109109

110+
void flushVariableWidthColumn(
111+
const std::vector<PartitionedVectorPtr>& partitionedVectors,
112+
const TypePtr& colType,
113+
const std::vector<uint32_t>& nonEmptyPartitions,
114+
const std::vector<IOBufOutputStream*>& outputStreams) const;
115+
110116
void flushSingleSimpleVector(
111117
const PartitionedVectorPtr& partitionedVector,
112118
const std::vector<IOBufOutputStream*>& outputStreams) const;
113119

120+
void flushSingleVariableWidthVector(
121+
const PartitionedVectorPtr& partitionedVector,
122+
const std::vector<IOBufOutputStream*>& outputStreams) const;
123+
114124
template <TypeKind kind>
115125
void flushSingleFlatVector(
116126
const PartitionedVectorPtr& partitionedVector,
117127
const std::vector<IOBufOutputStream*>& outputStreams) const;
118128

129+
void flushSingleVariableWidthFlatVector(
130+
const PartitionedVectorPtr& partitionedVector,
131+
const std::vector<IOBufOutputStream*>& outputStreams) const;
132+
119133
void flushHeader(
120134
std::string_view name,
121135
const std::vector<uint32_t>& nonEmptyPartitions,
@@ -141,6 +155,10 @@ class PrestoIterativePartitioningSerializer {
141155
const std::vector<uint32_t>& nonEmptyPartitions,
142156
const std::vector<IOBufOutputStream*>& outputStreams) const;
143157

158+
std::vector<int32_t> flushOffsets(
159+
const std::vector<PartitionedVectorPtr>& partitionedVectors,
160+
const std::vector<IOBufOutputStream*>& outputStreams) const;
161+
144162
RowTypePtr type_;
145163
uint32_t numPartitions_;
146164
SerdeOpts opts_;

‎velox/serializers/benchmarks/PrestoIterativePartitioningSerializerBenchmark.cpp‎

Lines changed: 39 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -44,21 +44,7 @@ class PrestoIterativePartitioningSerializerBenchmark
4444
}
4545

4646
/// Creates a flat vector of the given TypeKind with the given null ratio.
47-
VectorPtr makeColumn(vector_size_t size, TypeKind colKind, int32_t nullPct) {
48-
switch (colKind) {
49-
case TypeKind::BOOLEAN:
50-
return makeColumnOfType<bool>(size, nullPct);
51-
case TypeKind::INTEGER:
52-
return makeColumnOfType<int32_t>(size, nullPct);
53-
case TypeKind::BIGINT:
54-
return makeColumnOfType<int64_t>(size, nullPct);
55-
case TypeKind::HUGEINT:
56-
return makeColumnOfType<int128_t>(size, nullPct);
57-
default:
58-
VELOX_UNSUPPORTED(
59-
"Unsupported TypeKind: {}", TypeKindName::toName(colKind));
60-
}
61-
}
47+
VectorPtr makeColumn(vector_size_t size, TypeKind colKind, int32_t nullPct);
6248

6349
/// Creates a RowVector with numCols columns of the given TypeKind.
6450
RowVectorPtr makeInput(
@@ -96,6 +82,43 @@ class PrestoIterativePartitioningSerializerBenchmark
9682
}
9783
};
9884

85+
template <>
86+
VectorPtr PrestoIterativePartitioningSerializerBenchmark::makeColumnOfType<
87+
StringView>(vector_size_t size, int32_t nullPct) {
88+
auto valueAt = [](auto row) {
89+
return row % 2 == 0
90+
? StringView::makeInline(fmt::format("short{}", row))
91+
: StringView(fmt::format("long_string_value_{}_suffix", row).data());
92+
};
93+
94+
if (nullPct == 0) {
95+
return makeFlatVector<StringView>(size, valueAt);
96+
}
97+
return makeFlatVector<StringView>(
98+
size, valueAt, [nullPct](auto row) { return (row % 100) < nullPct; });
99+
}
100+
101+
VectorPtr PrestoIterativePartitioningSerializerBenchmark::makeColumn(
102+
vector_size_t size,
103+
TypeKind colKind,
104+
int32_t nullPct) {
105+
switch (colKind) {
106+
case TypeKind::BOOLEAN:
107+
return makeColumnOfType<bool>(size, nullPct);
108+
case TypeKind::INTEGER:
109+
return makeColumnOfType<int32_t>(size, nullPct);
110+
case TypeKind::BIGINT:
111+
return makeColumnOfType<int64_t>(size, nullPct);
112+
case TypeKind::HUGEINT:
113+
return makeColumnOfType<int128_t>(size, nullPct);
114+
case TypeKind::VARCHAR:
115+
case TypeKind::VARBINARY:
116+
return makeColumnOfType<StringView>(size, nullPct);
117+
default:
118+
VELOX_UNSUPPORTED(
119+
"Unsupported TypeKind: {}", TypeKindName::toName(colKind));
120+
}
121+
}
99122
} // namespace
100123

101124
/// Single benchmark function parameterized by (colKind, numCols, nullPct,
@@ -166,6 +189,7 @@ FLUSH_FOR_COLS(bool, BOOLEAN)
166189
FLUSH_FOR_COLS(int, INTEGER)
167190
FLUSH_FOR_COLS(bigint, BIGINT)
168191
FLUSH_FOR_COLS(ldec, HUGEINT)
192+
FLUSH_FOR_COLS(sv, VARCHAR)
169193
// clang-format on
170194

171195
int main(int argc, char** argv) {

0 commit comments

Comments
 (0)