Support only null type for json_extract (#8489)

close #8486
pingcap · Dec 11, 2023 · 171097e · 171097e
1 parent 3085414
commit 171097e
Show file tree

Hide file tree

Showing 3 changed files with 44 additions and 9 deletions.
diff --git a/dbms/src/Functions/FunctionsJson.h b/dbms/src/Functions/FunctionsJson.h
@@ -90,19 +90,14 @@ class FunctionJsonExtract : public IFunction
     {
         for (const auto & arg : arguments)
         {
-            if (const auto * nested_type = checkAndGetDataType<DataTypeNullable>(arg.get()))
+            if (!arg->onlyNull())
             {
-                if unlikely (!nested_type->getNestedType()->isStringOrFixedString())
+                const auto * nested_arg_type = removeNullable(arg).get();
+                if unlikely (!nested_arg_type->isStringOrFixedString())
                     throw Exception(
                         "Illegal type " + arg->getName() + " of argument of function " + getName(),
                         ErrorCodes::ILLEGAL_TYPE_OF_ARGUMENT);
             }
-            else if unlikely (!arg->isStringOrFixedString())
-            {
-                throw Exception(
-                    "Illegal type " + arg->getName() + " of argument of function " + getName(),
-                    ErrorCodes::ILLEGAL_TYPE_OF_ARGUMENT);
-            }
         }
         return makeNullable(std::make_shared<DataTypeString>());
     }
@@ -116,7 +111,12 @@ class FunctionJsonExtract : public IFunction
         bool const_json = false;
         const ColumnString * source_data_column_ptr;
         const ColumnNullable * source_nullable_column_ptr = nullptr;
-        if (const auto * const_nullable_col = checkAndGetColumnConst<ColumnNullable>(json_column))
+        if unlikely (json_column->onlyNull())
+        {
+            block.getByPosition(result).column = block.getByPosition(result).type->createColumnConst(rows, Null());
+            return;
+        }
+        else if (const auto * const_nullable_col = checkAndGetColumnConst<ColumnNullable>(json_column))
         {
             const_json = true;
             json_column = const_nullable_col->getDataColumnPtr().get();
@@ -175,6 +175,12 @@ class FunctionJsonExtract : public IFunction
         for (size_t i = 1; i < arguments_size; ++i)
         {
             const ColumnPtr column = block.getByPosition(arguments[i]).column;
+            if (column->onlyNull())
+            {
+                block.getByPosition(result).column = block.getByPosition(result).type->createColumnConst(rows, Null());
+                return;
+            }
+
             const auto * nested_column = static_cast<const ColumnConst *>(column.get())->getDataColumnPtr().get();
             StringRef path_str;
             if (const auto * nullable_string_path_col = checkAndGetColumn<ColumnNullable>(nested_column))

diff --git a/dbms/src/Functions/tests/gtest_json_extract.cpp b/dbms/src/Functions/tests/gtest_json_extract.cpp
@@ -137,6 +137,14 @@ try
     expect_null_vec = {1, 1, 1};
     checkResult(res.column, expect_null_vec, expect_string_vec);
 
+    /// JsonBinary only null
+    auto const_null_only_col = createOnlyNullColumnConst(3);
+    res = executeFunction(func_name, {const_null_only_col, path_col});
+    ASSERT_TRUE(res.column->size() == 3);
+    expect_string_vec = {"", "", ""};
+    expect_null_vec = {1, 1, 1};
+    checkResult(res.column, expect_null_vec, expect_string_vec);
+
     /// ColumnVector(non-null)
     auto non_null_str_col = ColumnString::create();
     non_null_str_col->insertData(reinterpret_cast<const char *>(bj2), sizeof(bj2) / sizeof(UInt8));
@@ -150,6 +158,13 @@ try
     expect_null_vec = {0, 0, 0};
     checkResult(res.column, expect_null_vec, expect_string_vec);
 
+    /// One of Paths is only null
+    res = executeFunction(func_name, {non_null_input_col, non_null_path_col, const_null_only_col});
+    ASSERT_TRUE(res.column->size() == 3);
+    expect_string_vec = {"", "", ""};
+    expect_null_vec = {1, 1, 1};
+    checkResult(res.column, expect_null_vec, expect_string_vec);
+
     /// ColumnConst(non-null)
     non_null_str_col = ColumnString::create();
     non_null_str_col->insertData(reinterpret_cast<const char *>(bj2), sizeof(bj2) / sizeof(UInt8));

diff --git a/tests/fullstack-test/expr/json_extract.test b/tests/fullstack-test/expr/json_extract.test
@@ -47,4 +47,18 @@ mysql> set tidb_enforce_mpp=1; select json_extract(d, '\$[0]', '\$[1]', '\$[2].a
 | [1, 2, "b"] |
 +-------------+
 
+mysql> set tidb_enforce_mpp=1; select json_extract(NULL, '\$[0]', '\$[1]', '\$[2].a') as col_a from test.t #NO_UNESCAPE
++-------+
+| col_a |
++-------+
+| NULL  |
++-------+
+
+mysql> set tidb_enforce_mpp=1; select json_extract(d, '\$[0]', NULL, '\$[2].a') as col_a from test.t #NO_UNESCAPE
++-------+
+| col_a |
++-------+
+| NULL  |
++-------+
+
 mysql> drop table if exists test.t