enable building of List Arrays from List Scalars

KatanaGraph · May 27, 2021 · 0dc4933 · 0dc4933
1 parent 3336ea5
commit 0dc4933
Showing 1 changed file with 42 additions and 23 deletions.
diff --git a/libsupport/src/ArrowVisitor.cpp b/libsupport/src/ArrowVisitor.cpp
@@ -12,6 +12,18 @@ struct ToArrayVisitor {
 
   using ReturnType = std::shared_ptr<arrow::Array>;
   using ResultType = katana::Result<ReturnType>;
+
+  template <typename BuilderType>
+  ResultType Finish(BuilderType* builder) {
+    std::shared_ptr<arrow::Array> array;
+    auto res = builder->Finish(&array);
+    if (!res.ok()) {
+      return KATANA_ERROR(
+          katana::ErrorCode::ArrowError, "arrow builder finish: {}", res);
+    }
+    return array;
+  }
+
   template <typename ArrowType, typename BuilderType>
   arrow::enable_if_t<
       arrow::is_number_type<ArrowType>::value ||
@@ -34,24 +46,20 @@ struct ToArrayVisitor {
         builder->UnsafeAppendNull();
       }
     }
-    std::shared_ptr<arrow::Array> array;
-    auto res = builder->Finish(&array);
-    if (!res.ok()) {
-      return KATANA_ERROR(
-          katana::ErrorCode::ArrowError, "arrow builder finish: {}", res);
-    }
-    return array;
+    return Finish(builder);
   }
+
   template <typename ArrowType, typename BuilderType>
   arrow::enable_if_string_like<ArrowType, ResultType> Call(
       BuilderType* builder) {
     using ScalarType = typename arrow::TypeTraits<ArrowType>::ScalarType;
-    // same as above, but use scalar.value->ToString()
+    // same as above, but with string_view and Append instead of UnsafeAppend
     for (const auto& scalar : scalars) {
       if (scalar->is_valid) {
         // ->value->ToString() works, scalar->ToString() yields "..."
         const ScalarType* typed_scalar = static_cast<ScalarType*>(scalar.get());
-        if (auto res = builder->Append(typed_scalar->value->ToString());
+        if (auto res = builder->Append(
+                (arrow::util::string_view)(*typed_scalar->value));
             !res.ok()) {
           return KATANA_ERROR(
               katana::ErrorCode::ArrowError, "arrow builder failed append: {}",
@@ -65,22 +73,33 @@ struct ToArrayVisitor {
         }
       }
     }
-    std::shared_ptr<arrow::Array> array;
-    auto res = builder->Finish(&array);
-    if (!res.ok()) {
-      return KATANA_ERROR(
-          katana::ErrorCode::ArrowError, "arrow builder finish: {}", res);
-    }
-    return array;
-  }
-  template <typename ArrowType, typename BuilderType>
-  katana::enable_if_list_type_patched<ArrowType, ResultType> Call(
-      BuilderType*) {
-    KATANA_LOG_FATAL("List visitors not yet supported");
+    return Finish(builder);
   }
+
   template <typename ArrowType, typename BuilderType>
-  arrow::enable_if_struct<ArrowType, ResultType> Call(BuilderType*) {
-    KATANA_LOG_FATAL("Struct visitors not yet supported");
+  std::enable_if_t<
+      katana::is_list_type_patched<ArrowType>::value ||
+          arrow::is_struct_type<ArrowType>::value,
+      ResultType>
+  Call(BuilderType* builder) {
+    using ScalarType = typename arrow::TypeTraits<ArrowType>::ScalarType;
+    // use a visitor to traverse more complex types
+    katana::AppendScalarToBuilder visitor(builder);
+    for (const auto& scalar : scalars) {
+      if (scalar->is_valid) {
+        const ScalarType* typed_scalar = static_cast<ScalarType*>(scalar.get());
+        if (auto res = visitor.Call<ArrowType>(*typed_scalar); !res) {
+          return res.error();
+        }
+      } else {
+        if (auto res = builder->AppendNull(); !res.ok()) {
+          return KATANA_ERROR(
+              katana::ErrorCode::ArrowError,
+              "arrow builder failed append null: {}", res);
+        }
+      }
+    }
+    return Finish(builder);
   }
 };