Address comments

facebookincubator · Mar 18, 2024 · 958e6b0 · 958e6b0
1 parent 0823abf
commit 958e6b0
Show file tree

Hide file tree

Showing 2 changed files with 40 additions and 24 deletions.
diff --git a/velox/vector/arrow/Bridge.cpp b/velox/vector/arrow/Bridge.cpp
@@ -926,6 +926,34 @@ void exportToArrowImpl(
   out.release = releaseArrowArray;
 }
 
+// Parses the velox decimal format from the given arrow format.
+// The input format string should be in the form "d:precision,scale<,bitWidth>".
+// bitWidth is not required and must be 128 if provided.
+TypePtr parseDecimalFormat(const char* format) {
+  try {
+    std::string::size_type sz;
+    // Parse "d:".
+    int precision = std::stoi(&format[2], &sz);
+    // Parse ",".
+    int idx = 2 + sz + 1;
+    int scale = std::stoi(&format[idx], &sz);
+    // If bitwidth is provided, check if it is equal to 128.
+    if (format[idx + sz] == ',') {
+      int bitWidth = std::stoi(&format[idx + sz + 1], &sz);
+      VELOX_USER_CHECK_EQ(
+          bitWidth,
+          128,
+          "Conversion failed for '{}'. Velox decimal does not support custom bitwidth.",
+          format);
+    }
+    return DECIMAL(precision, scale);
+  } catch (std::invalid_argument&) {
+    VELOX_USER_FAIL(
+        "Unable to convert '{}' ArrowSchema decimal format to Velox decimal",
+        format);
+  }
+}
+
 TypePtr importFromArrowImpl(
     const char* format,
     const ArrowSchema& arrowSchema) {
@@ -972,30 +1000,9 @@ TypePtr importFromArrowImpl(
       }
       break;
 
-    case 'd': { // decimal types.
-      try {
-        std::string::size_type sz;
-        // Parse "d:".
-        int precision = std::stoi(&format[2], &sz);
-        // Parse ",".
-        int idx = 2 + sz + 1;
-        int scale = std::stoi(&format[idx], &sz);
-        // If bitwidth is provided, check if it is equal to 128.
-        if (format[idx + sz] == ',') {
-          int bitWidth = std::stoi(&format[idx + sz + 1], &sz);
-          VELOX_USER_CHECK_EQ(
-              bitWidth,
-              128,
-              "Conversion failed for '{}'. Velox decimal does not support custom bitwidth.",
-              format);
-        }
-        return DECIMAL(precision, scale);
-      } catch (std::invalid_argument&) {
-        VELOX_USER_FAIL(
-            "Unable to convert '{}' ArrowSchema decimal format to Velox decimal",
-            format);
-      }
-    }
+    case 'd':
+      // decimal types.
+      return parseDecimalFormat(format);
 
     // Complex types.
     case '+': {

diff --git a/velox/vector/arrow/tests/ArrowBridgeSchemaTest.cpp b/velox/vector/arrow/tests/ArrowBridgeSchemaTest.cpp
@@ -378,6 +378,15 @@ TEST_F(ArrowBridgeSchemaImportTest, scalar) {
   VELOX_ASSERT_THROW(
       *testSchemaImport("d:10,4,"),
       "Unable to convert 'd:10,4,' ArrowSchema decimal format to Velox decimal");
+  VELOX_ASSERT_THROW(
+      *testSchemaImport("d:10"),
+      "Unable to convert 'd:10' ArrowSchema decimal format to Velox decimal");
+  VELOX_ASSERT_THROW(
+      *testSchemaImport("d:"),
+      "Unable to convert 'd:' ArrowSchema decimal format to Velox decimal");
+  VELOX_ASSERT_THROW(
+      *testSchemaImport("d:10,"),
+      "Unable to convert 'd:10,' ArrowSchema decimal format to Velox decimal");
 }
 
 TEST_F(ArrowBridgeSchemaImportTest, complexTypes) {