Complex type

kewang1024 · Apr 9, 2024 · c2554b1 · c2554b1
1 parent a11344a
commit c2554b1
Show file tree

Hide file tree

Showing 2 changed files with 56 additions and 0 deletions.
diff --git a/velox/vector/ComplexVector.cpp b/velox/vector/ComplexVector.cpp
@@ -448,6 +448,26 @@ uint64_t RowVector::estimateFlatSize() const {
   return total;
 }
 
+double RowVector::estimateRowSize() const {
+  uint64_t total;
+  for (const auto& child : children_) {
+    if (child) {
+      total += child->estimateRowSize();
+    }
+  }
+  return total;
+}
+
+uint64_t RowVector::estimateCompactSize() const {
+  uint64_t total = BaseVector::retainedSize();
+  for (const auto& child : children_) {
+    if (child) {
+      total += child->estimateCompactSize();
+    }
+  }
+  return total;
+}
+
 void RowVector::prepareForReuse() {
   BaseVector::prepareForReuse();
   for (auto& child : children_) {
@@ -979,6 +999,18 @@ uint64_t ArrayVector::estimateFlatSize() const {
       sizes_->capacity() + elements_->estimateFlatSize();
 }
 
+double ArrayVector::estimateRowSize() const {
+  if (BaseVector::length_ == 0) {
+    return 0;
+  }
+  auto avgCountPerRow = (rawOffsets_[BaseVector::length_-1] + rawSizes_[BaseVector::length_-1]) / BaseVector::length_;
+  return avgCountPerRow * elements_->estimateRowSize();
+}
+
+uint64_t ArrayVector::estimateCompactSize() const {
+  return BaseVector::retainedSize() + offsets_->capacity() + sizes_->capacity() + elements_->estimateCompactSize();
+}
+
 namespace {
 void zeroOutBuffer(BufferPtr buffer) {
   memset(buffer->asMutable<char>(), 0, buffer->size());
@@ -1283,6 +1315,18 @@ uint64_t MapVector::estimateFlatSize() const {
       values_->estimateFlatSize();
 }
 
+double MapVector::estimateRowSize() const {
+  if (BaseVector::length_ == 0) {
+    return 0;
+  }
+  auto avgCountPerRow = (rawOffsets_[BaseVector::length_-1] + rawSizes_[BaseVector::length_-1]) / BaseVector::length_;
+  return avgCountPerRow * (keys_->estimateRowSize() + values_->estimateRowSize());
+}
+
+uint64_t MapVector::estimateCompactSize() const {
+  return BaseVector::retainedSize() + offsets_->capacity() + sizes_->capacity() + keys_->estimateCompactSize() + values_->estimateCompactSize();
+}
+
 void MapVector::prepareForReuse() {
   BaseVector::prepareForReuse();
 

diff --git a/velox/vector/ComplexVector.h b/velox/vector/ComplexVector.h
@@ -180,6 +180,10 @@ class RowVector : public BaseVector {
 
   uint64_t estimateFlatSize() const override;
 
+  double estimateRowSize() const override;
+
+  uint64_t estimateCompactSize() const override;
+
   using BaseVector::toString;
 
   std::string toString(vector_size_t index) const override;
@@ -457,6 +461,10 @@ class ArrayVector : public ArrayVectorBase {
 
   uint64_t estimateFlatSize() const override;
 
+  double estimateRowSize() const override;
+
+  uint64_t estimateCompactSize() const override;
+
   using BaseVector::toString;
 
   std::string toString(vector_size_t index) const override;
@@ -586,6 +594,10 @@ class MapVector : public ArrayVectorBase {
 
   uint64_t estimateFlatSize() const override;
 
+  double estimateRowSize() const override;
+
+  uint64_t estimateCompactSize() const override;
+
   using BaseVector::toString;
 
   std::string toString(vector_size_t index) const override;