Merge pull request #2419 from kuzudb/arrow-fix

Add more types to arrow export
kuzudb · Nov 15, 2023 · 507e245 · 507e245
2 parents 80a8b59 + 1c5ef15
commit 507e245
Show file tree

Hide file tree

Showing 19 changed files with 491 additions and 161 deletions.
diff --git a/src/binder/binder.cpp b/src/binder/binder.cpp
@@ -90,7 +90,7 @@ std::unique_ptr<LogicalType> Binder::bindDataType(const std::string& dataType) {
     if (boundType.getLogicalTypeID() == LogicalTypeID::FIXED_LIST) {
         auto validNumericTypes = LogicalTypeUtils::getNumericalLogicalTypeIDs();
         auto childType = FixedListType::getChildType(&boundType);
-        auto numElementsInList = FixedListType::getNumElementsInList(&boundType);
+        auto numElementsInList = FixedListType::getNumValuesInList(&boundType);
         if (find(validNumericTypes.begin(), validNumericTypes.end(),
                 childType->getLogicalTypeID()) == validNumericTypes.end()) {
             throw BinderException("The child type of a fixed list must be a numeric type. Given: " +

diff --git a/src/c_api/data_type.cpp b/src/c_api/data_type.cpp
@@ -56,5 +56,5 @@ uint64_t kuzu_data_type_get_fixed_num_elements_in_list(kuzu_logical_type* data_t
     if (parent_type->getLogicalTypeID() != LogicalTypeID::FIXED_LIST) {
         return 0;
     }
-    return FixedListType::getNumElementsInList(static_cast<LogicalType*>(data_type->_data_type));
+    return FixedListType::getNumValuesInList(static_cast<LogicalType*>(data_type->_data_type));
 }
diff --git a/src/common/arrow/arrow_converter.cpp b/src/common/arrow/arrow_converter.cpp
@@ -3,7 +3,6 @@
 #include <cstring>
 
 #include "common/arrow/arrow_row_batch.h"
-#include "common/exception/internal.h"
 
 namespace kuzu {
 namespace common {
@@ -69,6 +68,9 @@ void ArrowConverter::setArrowFormat(
     case LogicalTypeID::BOOL: {
         child.format = "b";
     } break;
+    case LogicalTypeID::INT128: {
+        child.format = "d:38,0";
+    } break;
     case LogicalTypeID::INT64: {
         child.format = "l";
     } break;
@@ -78,9 +80,27 @@ void ArrowConverter::setArrowFormat(
     case LogicalTypeID::INT16: {
         child.format = "s";
     } break;
+    case LogicalTypeID::INT8: {
+        child.format = "c";
+    } break;
+    case LogicalTypeID::UINT64: {
+        child.format = "L";
+    } break;
+    case LogicalTypeID::UINT32: {
+        child.format = "I";
+    } break;
+    case LogicalTypeID::UINT16: {
+        child.format = "S";
+    } break;
+    case LogicalTypeID::UINT8: {
+        child.format = "C";
+    } break;
     case LogicalTypeID::DOUBLE: {
         child.format = "g";
     } break;
+    case LogicalTypeID::FLOAT: {
+        child.format = "f";
+    } break;
     case LogicalTypeID::DATE: {
         child.format = "tdD";
     } break;
@@ -105,16 +125,27 @@ void ArrowConverter::setArrowFormat(
         child.children[0]->name = "l";
         setArrowFormat(rootHolder, **child.children, *typeInfo.childrenTypesInfo[0]);
     } break;
+    case LogicalTypeID::FIXED_LIST: {
+        auto numValuesPerList = "+w:" + std::to_string(typeInfo.numValuesPerList);
+        child.format = copyName(rootHolder, numValuesPerList);
+        child.n_children = 1;
+        rootHolder.nestedChildren.emplace_back();
+        rootHolder.nestedChildren.back().resize(1);
+        rootHolder.nestedChildrenPtr.emplace_back();
+        rootHolder.nestedChildrenPtr.back().push_back(&rootHolder.nestedChildren.back()[0]);
+        initializeChild(rootHolder.nestedChildren.back()[0]);
+        child.children = &rootHolder.nestedChildrenPtr.back()[0];
+        child.children[0]->name = "l";
+        setArrowFormat(rootHolder, **child.children, *typeInfo.childrenTypesInfo[0]);
+    } break;
+    case LogicalTypeID::STRUCT:
     case LogicalTypeID::INTERNAL_ID:
     case LogicalTypeID::NODE:
     case LogicalTypeID::REL: {
         setArrowFormatForStruct(rootHolder, child, typeInfo);
     } break;
-        // LCOV_EXCL_START
     default:
-        throw InternalException(
-            "Unsupported Arrow type " + LogicalTypeUtils::toString(typeInfo.typeID));
-        // LCOV_EXCL_STOP
+        KU_UNREACHABLE;
     }
 }