Augment nimble_dump streams subcommand to include InMap stream indica…

…tor (#130)

Summary:

Add a new `--inmap_stream` option to the `streams` subcommand. When set, the output includes a new column, which indicates if each stream is an InMap stream for a FlatMap, or not.

Differential Revision: D67996331

dwio/nimble/tools/NimbleDump.cpp

@@ -136,6 +136,7 @@ int main(int argc, char* argv[]) {
                    options["no_header"].as<bool>(),
                    options["labels"].as<bool>(),
                    options["raw_size"].as<bool>(),
+                   options["inmap_stream"].as<bool>(),
                    getOptional<uint32_t>(options["stripe"]));
          },
          positionalArgs)
@@ -163,6 +164,10 @@ int main(int argc, char* argv[]) {
                 po::bool_switch()->default_value(false),
                 "Include stream labels. Lables provide a readable path from the "
                 "root node to the stream, as they appear in the schema tree."
+            )(
+                "inmap_stream,i",
+                po::bool_switch()->default_value(false),
+                "Include InMap stream indicator."
             );
   // clang-format on
 

dwio/nimble/tools/NimbleDumpLib.cpp

@@ -19,6 +19,7 @@
 #include <numeric>
 #include <ostream>
 #include <tuple>
+#include <unordered_set>
 #include <utility>
 
 #include "common/strings/Zstd.h"
@@ -444,6 +445,7 @@ void NimbleDumpLib::emitStreams(
     bool noHeader,
     bool showStreamLabels,
     bool showStreamRawSize,
+    bool showInMapStream,
     std::optional<uint32_t> stripeId) {
   auto tabletReader = std::make_shared<TabletReader>(*pool_, file_.get());
 
@@ -459,15 +461,31 @@ void NimbleDumpLib::emitStreams(
   if (showStreamLabels) {
     fields.push_back({"Stream Label", 16, Alignment::Left});
   }
+  if (showInMapStream) {
+    fields.push_back({"InMap Stream", 16, Alignment::Left});
+  }
   fields.push_back({"Type", 30, Alignment::Left});
 
-  TableFormatter formatter(ostream_, fields, noHeader);
+  TableFormatter formatter(ostream_, std::move(fields), noHeader);
 
   std::optional<StreamLabels> labels{};
   if (showStreamLabels) {
     VeloxReader reader{*pool_, tabletReader};
     labels.emplace(reader.schema());
   }
+  std::unordered_set<uint32_t> inMapStreams;
+  if (showInMapStream) {
+    VeloxReader reader{*pool_, tabletReader};
+    SchemaReader::traverseSchema(
+        reader.schema(), [&](auto /*level*/, const Type& type, auto /*info*/) {
+          if (type.kind() == Kind::FlatMap) {
+            auto& map = type.asFlatMap();
+            for (size_t i = 0; i < map.childrenCount(); ++i) {
+              inMapStreams.insert(map.inMapDescriptorAt(i).offset());
+            }
+          }
+        });
+  }
 
   traverseTablet(
       *pool_,
@@ -500,7 +518,10 @@ void NimbleDumpLib::emitStreams(
         }
         values.push_back(folly::to<std::string>(itemCount));
         if (showStreamLabels) {
-          auto it = values.emplace_back(labels->streamLabel(streamId));
+          values.emplace_back(labels->streamLabel(streamId));
+        }
+        if (showInMapStream) {
+          values.emplace_back(inMapStreams.contains(streamId) ? "T" : "F");
         }
         values.push_back(getStreamInputLabel(stream));
         formatter.writeRow(values);

dwio/nimble/tools/NimbleDumpLib.h

@@ -33,6 +33,7 @@ class NimbleDumpLib {
       bool noHeader,
       bool flatmapKeys,
       bool rawSize,
+      bool showInMapStream,
       std::optional<uint32_t> stripeId);
   void
   emitHistogram(bool topLevel, bool noHeader, std::optional<uint32_t> stripeId);