Initial commit

velox/functions/sparksql/Size.cpp

@@ -35,6 +35,18 @@ struct Size {
     legacySizeOfNull_ = config.sparkLegacySizeOfNull();
   }
 
+  template <typename TInput>
+  FOLLY_ALWAYS_INLINE void initialize(
+      const std::vector<TypePtr>& /*inputTypes*/,
+      const core::QueryConfig& config,
+      const TInput* /*input*/,
+      const bool* legacySizeOfNull) {
+    VELOX_USER_CHECK_NE(
+        legacySizeOfNull, nullptr, "Constant legacySizeOfNull is expected.");
+    // Respects the passed legacySizeOfNull.
+    legacySizeOfNull_ = *legacySizeOfNull;
+  }
+
   template <typename TInput>
   FOLLY_ALWAYS_INLINE bool callNullable(int32_t& out, const TInput* input) {
     if (input == nullptr) {
@@ -49,6 +61,14 @@ struct Size {
     return true;
   }
 
+  template <typename TInput>
+  FOLLY_ALWAYS_INLINE bool callNullable(
+      int32_t& out,
+      const TInput* input,
+      const bool* /*legacySizeOfNull*/) {
+    return callNullable(out, input);
+  }
+
  private:
   bool legacySizeOfNull_;
 };
@@ -57,6 +77,9 @@ struct Size {
 void registerSize(const std::string& prefix) {
   registerFunction<Size, int32_t, Array<Any>>({prefix});
   registerFunction<Size, int32_t, Map<Any, Any>>({prefix});
+  // Register with legacySizeOfNull.
+  registerFunction<Size, int32_t, Array<Any>, bool>({prefix});
+  registerFunction<Size, int32_t, Map<Any, Any>, bool>({prefix});
 }
 
 } // namespace facebook::velox::functions::sparksql

velox/functions/sparksql/tests/SizeTest.cpp

@@ -40,14 +40,22 @@ class SizeTest : public SparkFunctionBaseTest {
     }
   }
 
-  void testSizeLegacyNull(VectorPtr vector, vector_size_t numRows) {
+  void testSizeConfiguredLegacyNull(VectorPtr vector, vector_size_t numRows) {
     auto result =
         evaluate<SimpleVector<int32_t>>("size(c0)", makeRowVector({vector}));
     for (vector_size_t i = 0; i < numRows; ++i) {
       EXPECT_EQ(result->isNullAt(i), vector->isNullAt(i)) << "at " << i;
     }
   }
 
+  void testSizePassedLegacyNull(VectorPtr vector, vector_size_t numRows) {
+    auto result = evaluate<SimpleVector<int32_t>>(
+        "size(c0, false)", makeRowVector({vector}));
+    for (vector_size_t i = 0; i < numRows; ++i) {
+      EXPECT_EQ(result->isNullAt(i), vector->isNullAt(i)) << "at " << i;
+    }
+  }
+
   void setConfig(std::string configStr, bool value) {
     execCtx_.queryCtx()->testingOverrideConfigUnsafe({
         {configStr, std::to_string(value)},
@@ -75,16 +83,31 @@ TEST_F(SizeTest, sizetest) {
   testSize(mapVector, numRows);
 }
 
-// Ensure that out if set to -1 if SparkLegacySizeOfNull is specified.
-TEST_F(SizeTest, legacySizeOfNull) {
+// Ensure that out is set to null for null input if configured
+// SparkLegacySizeOfNull is false.
+TEST_F(SizeTest, configuredLegacySizeOfNull) {
   vector_size_t numRows = 100;
   setConfig(core::QueryConfig::kSparkLegacySizeOfNull, false);
   auto arrayVector =
       makeArrayVector<int64_t>(numRows, sizeAt, valueAt, nullEvery(1));
-  testSizeLegacyNull(arrayVector, numRows);
+  testSizeConfiguredLegacyNull(arrayVector, numRows);
+  auto mapVector = makeMapVector<int64_t, int64_t>(
+      numRows, sizeAt, valueAt, valueAt, nullEvery(1));
+  testSizeConfiguredLegacyNull(mapVector, numRows);
+}
+
+// Ensure that out is set to null for null input if passed legacySizeOfNull is
+// false, regardless of the configuration.
+TEST_F(SizeTest, passedLegacySizeOfNull) {
+  vector_size_t numRows = 100;
+  // Dismiss the configuration.
+  setConfig(core::QueryConfig::kSparkLegacySizeOfNull, true);
+  auto arrayVector =
+      makeArrayVector<int64_t>(numRows, sizeAt, valueAt, nullEvery(1));
+  testSizePassedLegacyNull(arrayVector, numRows);
   auto mapVector = makeMapVector<int64_t, int64_t>(
       numRows, sizeAt, valueAt, valueAt, nullEvery(1));
-  testSizeLegacyNull(mapVector, numRows);
+  testSizePassedLegacyNull(mapVector, numRows);
 }
 
 } // namespace facebook::velox::functions::sparksql::test