spark split only support integer type limit

gaoyangxiaozhu · Jul 2, 2024 · 60700f2 · 60700f2
1 parent be3d8df
commit 60700f2
Showing 1 changed file with 8 additions and 8 deletions.
diff --git a/velox/functions/sparksql/SplitFunctions.cpp b/velox/functions/sparksql/SplitFunctions.cpp
@@ -41,16 +41,16 @@ class Split final : public exec::VectorFunction {
     BaseVector::ensureWritable(rows, ARRAY(VARCHAR()), context.pool(), result);
     exec::VectorWriter<Array<Varchar>> resultWriter;
     resultWriter.init(*result->as<ArrayVector>());
-    int64_t limit = std::numeric_limits<int64_t>::max();
+    int32_t limit = std::numeric_limits<int32_t>::max();
     // Optimization for the (flat, const, const) case.
     if (strings->isIdentityMapping() and delims->isConstantMapping() and
         (noLimit or limits->isConstantMapping())) {
       const auto* rawStrings = strings->data<StringView>();
       const auto delim = delims->valueAt<StringView>(0);
       if (!noLimit) {
-        limit = limits->valueAt<int64_t>(0);
+        limit = limits->valueAt<int32_t>(0);
         if (limit <= 0) {
-          limit = std::numeric_limits<int64_t>::max();
+          limit = std::numeric_limits<int32_t>::max();
         }
       }
       rows.applyToSelected([&](vector_size_t row) {
@@ -66,9 +66,9 @@ class Split final : public exec::VectorFunction {
       rows.applyToSelected([&](vector_size_t row) {
         const auto delim = delims->valueAt<StringView>(row);
         if (!noLimit) {
-          limit = limits->valueAt<int64_t>(row);
+          limit = limits->valueAt<int32_t>(row);
           if (limit <= 0) {
-            limit = std::numeric_limits<int64_t>::max();
+            limit = std::numeric_limits<int32_t>::max();
           }
         }
         if (delim.size() == 0) {
@@ -103,7 +103,7 @@ class Split final : public exec::VectorFunction {
   // "bc"].
   void splitEmptyDelimer(
       const StringView current,
-      int64_t limit,
+      int32_t limit,
       vector_size_t row,
       exec::VectorWriter<Array<Varchar>>& resultWriter) const {
     resultWriter.setOffset(row);
@@ -128,7 +128,7 @@ class Split final : public exec::VectorFunction {
   void splitInner(
       StringView input,
       const StringView delim,
-      int64_t limit,
+      int32_t limit,
       vector_size_t row,
       exec::VectorWriter<Array<Varchar>>& resultWriter) const {
     // Add new array (for the new row) to our array vector.
@@ -201,7 +201,7 @@ std::vector<std::shared_ptr<exec::FunctionSignature>> signatures() {
                               .returnType("array(varchar)")
                               .argumentType("varchar")
                               .argumentType("varchar")
-                              .argumentType("bigint")
+                              .argumentType("integer")
                               .build());
   return signatures;
 }