Pretty-print REE arrays

Lead-authored-by: Tobias Zagorni <tobias@zagorni.eu> Co-authored-by: Felipe Oliveira Carvalho <felipekde@gmail.com>
apache · Jan 23, 2023 · 5f9cee8 · 5f9cee8
1 parent 43e3688
commit 5f9cee8
Show file tree

Hide file tree

Showing 2 changed files with 85 additions and 1 deletion.
diff --git a/cpp/src/arrow/array/array_run_end_test.cc b/cpp/src/arrow/array/array_run_end_test.cc
@@ -27,6 +27,7 @@
 #include "arrow/array/builder_run_end.h"
 #include "arrow/array/concatenate.h"
 #include "arrow/chunked_array.h"
+#include "arrow/pretty_print.h"
 #include "arrow/scalar.h"
 #include "arrow/status.h"
 #include "arrow/testing/builder.h"
@@ -463,6 +464,77 @@ TEST_P(TestRunEndEncodedArray, Concatenate) {
       Concatenate({int32_array, string_array}, default_memory_pool()));
 }
 
+TEST_P(TestRunEndEncodedArray, Printing) {
+  ASSERT_OK_AND_ASSIGN(auto int_array,
+                       RunEndEncodedArray::Make(size_values, int32_values, 30));
+  std::stringstream ss;
+  ASSERT_OK(PrettyPrint(*int_array, {}, &ss));
+  ASSERT_EQ(ss.str(),
+            "\n"
+            "-- run ends array (offset: 0, logical length: 30)\n"
+            "  [\n"
+            "    10,\n"
+            "    20,\n"
+            "    30\n"
+            "  ]\n"
+            "-- values:\n"
+            "  [\n"
+            "    10,\n"
+            "    20,\n"
+            "    30\n"
+            "  ]");
+
+  ASSERT_OK_AND_ASSIGN(auto string_array,
+                       RunEndEncodedArray::Make(size_values, string_values, 30));
+  ss = {};
+  ASSERT_OK(PrettyPrint(*string_array, {}, &ss));
+  ASSERT_EQ(ss.str(),
+            "\n"
+            "-- run ends array (offset: 0, logical length: 30)\n"
+            "  [\n"
+            "    10,\n"
+            "    20,\n"
+            "    30\n"
+            "  ]\n"
+            "-- values:\n"
+            "  [\n"
+            "    \"Hello\",\n"
+            "    \"World\",\n"
+            "    null\n"
+            "  ]");
+
+  auto sliced_array = string_array->Slice(15, 6);
+  ss = {};
+  ASSERT_OK(PrettyPrint(*sliced_array, {}, &ss));
+  ASSERT_EQ(ss.str(),
+            "\n"
+            "-- run ends array (offset: 15, logical length: 6)\n"
+            "  [\n"
+            "    10,\n"
+            "    20,\n"
+            "    30\n"
+            "  ]\n"
+            "-- values:\n"
+            "  [\n"
+            "    \"Hello\",\n"
+            "    \"World\",\n"
+            "    null\n"
+            "  ]");
+
+  ASSERT_OK_AND_ASSIGN(auto empty_array,
+                       RunEndEncodedArray::Make(ArrayFromJSON(run_ends_type, "[]"),
+                                                ArrayFromJSON(binary(), "[]"), 0));
+
+  ss = {};
+  ASSERT_OK(PrettyPrint(*empty_array, {}, &ss));
+  ASSERT_EQ(ss.str(),
+            "\n"
+            "-- run ends array (offset: 0, logical length: 0)\n"
+            "  []\n"
+            "-- values:\n"
+            "  []");
+}
+
 }  // anonymous namespace
 
 INSTANTIATE_TEST_SUITE_P(EncodedArrayTests, TestRunEndEncodedArray,

diff --git a/cpp/src/arrow/pretty_print.cc b/cpp/src/arrow/pretty_print.cc
@@ -378,7 +378,19 @@ class ArrayPrinter : public PrettyPrinter {
   }
 
   Status Visit(const RunEndEncodedArray& array) {
-    return Status::NotImplemented("printing run-end encoded array");
+    Newline();
+    Indent();
+    Write("-- run ends array (offset: ");
+    Write(std::to_string(array.offset()));
+    Write(", logical length: ");
+    Write(std::to_string(array.length()));
+    Write(")\n");
+    RETURN_NOT_OK(PrettyPrint(*array.run_ends_array(), ChildOptions(true), sink_));
+
+    Newline();
+    Indent();
+    Write("-- values:\n");
+    return PrettyPrint(*array.values_array(), ChildOptions(true), sink_);
   }
 
   Status Print(const Array& array) {