Functions: fix json valid and json extract in conditions

dbms/src/Flash/Coprocessor/DAGExpressionAnalyzer.cpp

@@ -49,6 +49,8 @@
 #include <tipb/executor.pb.h>
 #include <tipb/expression.pb.h>
 
+#include <ext/scope_guard.h>
+
 namespace DB
 {
 namespace ErrorCodes
@@ -1001,6 +1003,13 @@ String DAGExpressionAnalyzer::buildFilterColumn(
     const google::protobuf::RepeatedPtrField<tipb::Expr> & conditions,
     bool null_as_false)
 {
+    building_filter_conditions = true;
+    json_valid_guarded_exprs.clear();
+    SCOPE_EXIT({
+        building_filter_conditions = false;
+        json_valid_guarded_exprs.clear();
+    });
+
     String filter_column_name;
     if (conditions.size() == 1)
     {
@@ -1021,7 +1030,12 @@ String DAGExpressionAnalyzer::buildFilterColumn(
     {
         Names arg_names;
         for (const auto & condition : conditions)
+        {
+            auto guards_before_condition = json_valid_guarded_exprs;
             arg_names.push_back(getActions(condition, actions, true));
+            json_valid_guarded_exprs = std::move(guards_before_condition);
+            recordJsonValidGuards(condition);
+        }
         // connect all the conditions by logical and
         // two_value_and treats null as false inside the `two_value_and` function, so the output column
         // will always be UInt8 type, which can save the merge step in FilterDescription
@@ -1032,6 +1046,30 @@ String DAGExpressionAnalyzer::buildFilterColumn(
     return filter_column_name;
 }
 
+void DAGExpressionAnalyzer::recordJsonValidGuards(const tipb::Expr & expr)
+{
+    if (!building_filter_conditions || !isScalarFunctionExpr(expr))
+        return;
+
+    if (expr.sig() == tipb::ScalarFuncSig::JsonValidStringSig && expr.children_size() == 1)
+    {
+        json_valid_guarded_exprs.emplace(exprToString(expr.children(0), getCurrentInputColumns()));
+        return;
+    }
+
+    if (expr.sig() == tipb::ScalarFuncSig::LogicalAnd)
+    {
+        for (const auto & child : expr.children())
+            recordJsonValidGuards(child);
+    }
+}
+
+bool DAGExpressionAnalyzer::isJsonValidGuarded(const tipb::Expr & expr) const
+{
+    return building_filter_conditions
+        && json_valid_guarded_exprs.contains(exprToString(expr, getCurrentInputColumns()));
+}
+
 std::tuple<ExpressionActionsPtr, String, ExpressionActionsPtr> DAGExpressionAnalyzer::buildPushDownFilter(
     const google::protobuf::RepeatedPtrField<tipb::Expr> & conditions,
     bool null_as_false)

dbms/src/Flash/Coprocessor/DAGExpressionAnalyzer.h

@@ -28,6 +28,8 @@
 #include <Interpreters/WindowDescription.h>
 #include <Storages/KVStore/TMTStorages.h>
 
+#include <unordered_set>
+
 namespace DB
 {
 class Set;
@@ -318,11 +320,17 @@ class DAGExpressionAnalyzer : private boost::noncopyable
         const std::vector<tipb::FieldType> & require_schema,
         const std::vector<Int32> & output_offsets) const;
 
+    void recordJsonValidGuards(const tipb::Expr & expr);
+    bool isJsonValidGuarded(const tipb::Expr & expr) const;
+
     // all columns from table scan
     NamesAndTypes source_columns;
     DAGPreparedSets prepared_sets;
     const Context & context;
 
+    bool building_filter_conditions = false;
+    std::unordered_set<String> json_valid_guarded_exprs;
+
     friend class DAGExpressionAnalyzerHelper;
 };
 

dbms/src/Flash/Coprocessor/DAGExpressionAnalyzerHelper.cpp

@@ -202,13 +202,20 @@ String DAGExpressionAnalyzerHelper::buildLogicalFunction(
     const ExpressionActionsPtr & actions)
 {
     const String & func_name = getFunctionName(expr);
+    auto guards_before_function = analyzer->json_valid_guarded_exprs;
     Names argument_names;
     for (const auto & child : expr.children())
     {
+        auto guards_before_child = analyzer->json_valid_guarded_exprs;
         String name = analyzer->getActions(child, actions, true);
         argument_names.push_back(name);
+        analyzer->json_valid_guarded_exprs = std::move(guards_before_child);
+        if (func_name == "and" || func_name == "two_value_and")
+            analyzer->recordJsonValidGuards(child);
     }
-    return analyzer->applyFunction(func_name, argument_names, actions, getCollatorFromExpr(expr));
+    String result = analyzer->applyFunction(func_name, argument_names, actions, getCollatorFromExpr(expr));
+    analyzer->json_valid_guarded_exprs = std::move(guards_before_function);
+    return result;
 }
 
 // left(str,len) = substrUTF8(str,1,len)
@@ -299,7 +306,12 @@ String DAGExpressionAnalyzerHelper::buildSingleParamJsonRelatedFunctions(
     const auto & input_expr = expr.children(0);
     String arg = analyzer->getActions(input_expr, actions);
     const auto & collator = getCollatorFromExpr(expr);
+    const bool ignore_invalid_json
+        = func_name == FunctionCastStringAsJson::name && analyzer->isJsonValidGuarded(input_expr);
     String result_name = genFuncString(func_name, {arg}, {collator}, {&input_expr.field_type(), &expr.field_type()});
+    // Guarded and strict casts can coexist in different logical branches and must not share an action.
+    if (ignore_invalid_json)
+        result_name += "_json_valid_guarded";
     if (actions->getSampleBlock().has(result_name))
         return result_name;
 
@@ -318,6 +330,7 @@ String DAGExpressionAnalyzerHelper::buildSingleParamJsonRelatedFunctions(
         {
             function_cast_string_as_json->setInputTiDBFieldType(input_expr.field_type());
             function_cast_string_as_json->setOutputTiDBFieldType(expr.field_type());
+            function_cast_string_as_json->setIgnoreInvalidJson(ignore_invalid_json);
         }
         else if (auto * function_cast_time_as_json = dynamic_cast<FunctionCastTimeAsJson *>(function_impl);
                  function_cast_time_as_json)

dbms/src/Functions/FunctionsJson.h

@@ -1473,6 +1473,7 @@ class FunctionCastStringAsJson : public IFunction
 
     void setInputTiDBFieldType(const tipb::FieldType & tidb_tp_) { input_tidb_tp = tidb_tp_; }
     void setOutputTiDBFieldType(const tipb::FieldType & tidb_tp_) { output_tidb_tp = tidb_tp_; }
+    void setIgnoreInvalidJson(bool value) { ignore_invalid_json = value; }
     void setCollator(const TiDB::TiDBCollatorPtr & collator_) override { collator = collator_; }
 
     DataTypePtr getReturnTypeImpl(const DataTypes & arguments) const override
@@ -1572,11 +1573,18 @@ class FunctionCastStringAsJson : public IFunction
                     offsets_to,
                     input_source,
                     column_nullable.getNullMapData(),
-                    block.rows());
+                    block.rows(),
+                    ignore_invalid_json);
             }
             else
             {
-                doExecuteForParsingJson<false>(data_to, offsets_to, input_source, {}, block.rows());
+                doExecuteForParsingJson<false>(
+                    data_to,
+                    offsets_to,
+                    input_source,
+                    {},
+                    block.rows(),
+                    ignore_invalid_json);
             }
         }
         else
@@ -1697,7 +1705,8 @@ class FunctionCastStringAsJson : public IFunction
         ColumnString::Offsets & offsets_to,
         const std::unique_ptr<IStringSource> & data_from,
         const NullMap & null_map_from,
-        size_t size)
+        size_t size,
+        bool ignore_invalid_json)
     {
         // json_type + size of data_from.
         size_t reserve_size = size + data_from->getSizeForReserve();
@@ -1718,16 +1727,32 @@ class FunctionCastStringAsJson : public IFunction
 
             const auto & slice = data_from->getWhole();
             if (unlikely(slice.size == 0))
-                throw Exception("Invalid JSON text: The document is empty.");
+            {
+                if (!ignore_invalid_json)
+                    throw Exception("Invalid JSON text: The document is empty.");
+                JsonBinary::appendNull(write_buffer);
+                writeChar(0, write_buffer);
+                offsets_to[i] = write_buffer.count();
+                data_from->next();
+                continue;
+            }
 
             const auto & json_elem = parser.parse(slice.data, slice.size);
             if (unlikely(json_elem.error()))
             {
-                throw Exception(fmt::format(
-                    "Invalid JSON text: The document root must not be followed by other values, details: {}",
-                    simdjson::error_message(json_elem.error())));
+                if (!ignore_invalid_json || checkJsonValid(reinterpret_cast<const char *>(slice.data), slice.size))
+                {
+                    throw Exception(fmt::format(
+                        "Invalid JSON text: The document root must not be followed by other values, details: {}",
+                        simdjson::error_message(json_elem.error())));
+                }
+                // Keep vectorized evaluation alive until the matching JSON_VALID conjunct filters this row.
+                JsonBinary::appendNull(write_buffer);
+            }
+            else
+            {
+                JsonBinary::appendSIMDJsonElem(write_buffer, json_elem.value_unsafe());
             }
-            JsonBinary::appendSIMDJsonElem(write_buffer, json_elem.value_unsafe());
 
             writeChar(0, write_buffer);
             offsets_to[i] = write_buffer.count();
@@ -1758,6 +1783,7 @@ class FunctionCastStringAsJson : public IFunction
     std::optional<tipb::FieldType> input_tidb_tp;
     std::optional<tipb::FieldType> output_tidb_tp;
     TiDB::TiDBCollatorPtr collator = nullptr;
+    bool ignore_invalid_json = false;
 };
 
 class FunctionCastTimeAsJson : public IFunction

dbms/src/Functions/tests/gtest_json_valid.cpp

@@ -13,9 +13,14 @@
 // limitations under the License.
 
 #include <Columns/ColumnNullable.h>
+#include <Flash/Coprocessor/DAGCodec.h>
+#include <Flash/Coprocessor/DAGExpressionAnalyzer.h>
+#include <IO/Buffer/WriteBufferFromString.h>
+#include <Interpreters/ExpressionActions.h>
 #include <TestUtils/FunctionTestUtils.h>
 #include <TestUtils/TiFlashTestBasic.h>
 #include <TiDB/Decode/JsonBinary.h>
+#include <TiDB/Schema/TiDB.h>
 
 #include <string>
 #include <vector>
@@ -87,4 +92,114 @@ try
 }
 CATCH
 
+TEST_F(TestJsonValid, GuardStringToJsonParsingInFilter)
+try
+{
+    getDAGContext().log = Logger::get("TestJsonValid");
+
+    auto make_field_type = [](Int32 tp, UInt32 flag = 0) {
+        tipb::FieldType field_type;
+        field_type.set_tp(tp);
+        field_type.set_flag(flag);
+        return field_type;
+    };
+    auto make_column_ref = [&] {
+        tipb::Expr expr;
+        expr.set_tp(tipb::ExprType::ColumnRef);
+        WriteBufferFromOwnString ss;
+        encodeDAGInt64(0, ss);
+        expr.set_val(ss.releaseStr());
+        *expr.mutable_field_type() = make_field_type(TiDB::TypeString);
+        return expr;
+    };
+    auto make_scalar = [](tipb::ScalarFuncSig sig, const tipb::FieldType & field_type) {
+        tipb::Expr expr;
+        expr.set_tp(tipb::ExprType::ScalarFunc);
+        expr.set_sig(sig);
+        *expr.mutable_field_type() = field_type;
+        return expr;
+    };
+
+    const auto column_ref = make_column_ref();
+    auto json_valid = make_scalar(
+        tipb::ScalarFuncSig::JsonValidStringSig,
+        make_field_type(TiDB::TypeLongLong, TiDB::ColumnFlagIsBooleanFlag));
+    *json_valid.add_children() = column_ref;
+
+    auto cast_json = make_scalar(
+        tipb::ScalarFuncSig::CastStringAsJson,
+        make_field_type(TiDB::TypeJSON, TiDB::ColumnFlagParseToJSON));
+    *cast_json.add_children() = column_ref;
+
+    auto is_null = make_scalar(
+        tipb::ScalarFuncSig::StringIsNull,
+        make_field_type(TiDB::TypeLongLong, TiDB::ColumnFlagIsBooleanFlag));
+    *is_null.add_children() = cast_json;
+
+    auto is_not_null = make_scalar(
+        tipb::ScalarFuncSig::UnaryNotInt,
+        make_field_type(TiDB::TypeLongLong, TiDB::ColumnFlagIsBooleanFlag));
+    *is_not_null.add_children() = is_null;
+
+    auto execute_filter = [&](const google::protobuf::RepeatedPtrField<tipb::Expr> & conditions) {
+        Block block({createColumn<String>({"", "invalid json", R"({"a": 1})"}, "json")});
+        auto actions = std::make_shared<ExpressionActions>(block.getColumnsWithTypeAndName());
+        DAGExpressionAnalyzer analyzer(block, *context);
+        const auto filter_column = analyzer.buildFilterColumn(actions, conditions, true);
+        actions->execute(block);
+        return block.getByName(filter_column);
+    };
+
+    google::protobuf::RepeatedPtrField<tipb::Expr> guarded_conditions;
+    *guarded_conditions.Add() = json_valid;
+    *guarded_conditions.Add() = is_not_null;
+    ASSERT_COLUMN_EQ(createColumn<UInt8>({0, 0, 1}), execute_filter(guarded_conditions));
+
+    google::protobuf::RepeatedPtrField<tipb::Expr> reversed_conditions;
+    *reversed_conditions.Add() = is_not_null;
+    *reversed_conditions.Add() = json_valid;
+    ASSERT_THROW(execute_filter(reversed_conditions), Exception);
+
+    google::protobuf::RepeatedPtrField<tipb::Expr> unguarded_conditions;
+    *unguarded_conditions.Add() = is_not_null;
+    ASSERT_THROW(execute_filter(unguarded_conditions), Exception);
+
+    auto nested_and = make_scalar(
+        tipb::ScalarFuncSig::LogicalAnd,
+        make_field_type(TiDB::TypeLongLong, TiDB::ColumnFlagIsBooleanFlag));
+    *nested_and.add_children() = json_valid;
+    *nested_and.add_children() = json_valid;
+    auto guarded_nested_and = make_scalar(
+        tipb::ScalarFuncSig::LogicalAnd,
+        make_field_type(TiDB::TypeLongLong, TiDB::ColumnFlagIsBooleanFlag));
+    *guarded_nested_and.add_children() = nested_and;
+    *guarded_nested_and.add_children() = is_not_null;
+    google::protobuf::RepeatedPtrField<tipb::Expr> nested_conditions;
+    *nested_conditions.Add() = guarded_nested_and;
+    ASSERT_COLUMN_EQ(createColumn<UInt8>({0, 0, 1}), execute_filter(nested_conditions));
+
+    auto guarded_and = make_scalar(
+        tipb::ScalarFuncSig::LogicalAnd,
+        make_field_type(TiDB::TypeLongLong, TiDB::ColumnFlagIsBooleanFlag));
+    *guarded_and.add_children() = json_valid;
+    *guarded_and.add_children() = is_not_null;
+    auto unguarded_or = make_scalar(
+        tipb::ScalarFuncSig::LogicalOr,
+        make_field_type(TiDB::TypeLongLong, TiDB::ColumnFlagIsBooleanFlag));
+    *unguarded_or.add_children() = guarded_and;
+    *unguarded_or.add_children() = is_not_null;
+    google::protobuf::RepeatedPtrField<tipb::Expr> or_conditions;
+    *or_conditions.Add() = unguarded_or;
+    ASSERT_THROW(execute_filter(or_conditions), Exception);
+
+    auto non_and_wrapper
+        = make_scalar(tipb::ScalarFuncSig::EQInt, make_field_type(TiDB::TypeLongLong, TiDB::ColumnFlagIsBooleanFlag));
+    *non_and_wrapper.add_children() = guarded_and;
+    *non_and_wrapper.add_children() = is_not_null;
+    google::protobuf::RepeatedPtrField<tipb::Expr> wrapped_conditions;
+    *wrapped_conditions.Add() = non_and_wrapper;
+    ASSERT_THROW(execute_filter(wrapped_conditions), Exception);
+}
+CATCH
+
 } // namespace DB::tests

tests/fullstack-test/expr/json_valid.test

@@ -28,5 +28,19 @@ mysql> set tidb_allow_mpp=1;set tidb_enforce_mpp=1; set tidb_isolation_read_engi
 |                    1 |                      0 |                      0 |
 +----------------------+------------------------+------------------------+
 
+mysql> set tidb_opt_enable_late_materialization=0; select col_string->>'$.a' as extracted from test.t where json_valid(col_string) and col_string->>'$.a' is not null order by extracted;
++-----------+
+| extracted |
++-----------+
+| b         |
++-----------+
+
+mysql> set tidb_opt_enable_late_materialization=1; select col_string->>'$.a' as extracted from test.t where json_valid(col_string) and col_string->>'$.a' is not null order by extracted;
++-----------+
+| extracted |
++-----------+
+| b         |
++-----------+
+
 # Clean up.
 mysql> drop table if exists test.t;