Extract Null Count from Nimble Files (#177)

Summary:

Rollback Plan:

Differential Revision: D75928862

Author: macvincent

dwio/nimble/velox/RawSizeContext.h

@@ -57,9 +57,27 @@ class RawSizeContext {
     return columnSizes_.size();
   }
 
+  void appendNullCount(uint64_t nullCount) {
+    columnNullCounts_.push_back(nullCount);
+  }
+
+  uint64_t nullsAt(uint64_t columnIndex) const {
+    NIMBLE_ASSERT(
+        columnIndex < columnNullCounts_.size(),
+        fmt::format(
+            "Column index {} is out of range. Total number of columns is {}",
+            columnIndex,
+            columnNullCounts_.size()));
+    return columnNullCounts_.at(columnIndex);
+  }
+
+  // Number of nulls in last visited node
+  uint64_t nulls{0};
+
  private:
   DecodedVectorManager decodedVectorManager_;
   std::vector<uint64_t> columnSizes_;
+  std::vector<uint64_t> columnNullCounts_;
 };
 
 } // namespace facebook::nimble

dwio/nimble/velox/RawSizeUtils.cpp

@@ -56,6 +56,7 @@ uint64_t getRawSizeFromFixedWidthVector(
         }
       }
 
+      context.nulls = nullCount;
       return ((ranges.size() - nullCount) * sizeof(T)) +
           (nullCount * NULL_SIZE);
     }
@@ -67,6 +68,7 @@ uint64_t getRawSizeFromFixedWidthVector(
           encoding,
           vector->typeKind());
 
+      context.nulls = constVector->mayHaveNulls() ? ranges.size() : 0;
       return constVector->mayHaveNulls() ? ranges.size() * NULL_SIZE
                                          : ranges.size() * sizeof(T);
     }
@@ -92,6 +94,7 @@ uint64_t getRawSizeFromFixedWidthVector(
         }
       }
 
+      context.nulls = nullCount;
       return ((ranges.size() - nullCount) * sizeof(T)) +
           (nullCount * NULL_SIZE);
     }
@@ -132,6 +135,7 @@ uint64_t getRawSizeFromStringVector(
         }
       }
 
+      context.nulls = nullCount;
       return rawSize + (nullCount * NULL_SIZE);
     }
     case velox::VectorEncoding::Simple::CONSTANT: {
@@ -143,6 +147,7 @@ uint64_t getRawSizeFromStringVector(
           encoding,
           vector->typeKind());
 
+      context.nulls = constVector->mayHaveNulls() ? ranges.size() : 0;
       return constVector->mayHaveNulls()
           ? ranges.size() * NULL_SIZE
           : ranges.size() * constVector->value().size();
@@ -179,6 +184,7 @@ uint64_t getRawSizeFromStringVector(
         }
       }
 
+      context.nulls = nullCount;
       return rawSize + (nullCount * NULL_SIZE);
     }
     default: {
@@ -224,7 +230,7 @@ uint64_t getRawSizeFromConstantComplexVector(
   } else {
     rawSize = getRawSizeFromVector(valueVector, childRanges, context);
   }
-
+  context.nulls = constantVector->mayHaveNulls() ? ranges.size() : 0;
   return rawSize * ranges.size();
 }
 
@@ -333,6 +339,7 @@ uint64_t getRawSizeFromArrayVector(
         getRawSizeFromVector(arrayVector->elements(), childRanges, context);
   }
 
+  context.nulls = nullCount;
   if (nullCount) {
     rawSize += nullCount * NULL_SIZE;
   }
@@ -446,6 +453,7 @@ uint64_t getRawSizeFromMapVector(
         getRawSizeFromVector(mapVector->mapValues(), childRanges, context);
   }
 
+  context.nulls = nullCount;
   if (nullCount) {
     rawSize += nullCount * NULL_SIZE;
   }
@@ -549,10 +557,12 @@ uint64_t getRawSizeFromRowVector(
       rawSize += childRawSize;
       if (topLevel) {
         context.appendSize(childRawSize);
+        context.appendNullCount(context.nulls);
       }
     }
   }
 
+  context.nulls = nullCount;
   if (nullCount) {
     rawSize += nullCount * NULL_SIZE;
   }