YQ RD support strings parsing in Json columns (#11341)

Author: GrigoriyPA

ydb/core/fq/libs/row_dispatcher/json_parser.cpp

@@ -139,6 +139,10 @@ class TColumnParser {
     }
 
     static TParser GetJsonValueParser(NYql::NUdf::EDataSlot dataSlot, bool optional) {
+        if (dataSlot == NYql::NUdf::EDataSlot::Json) {
+            return GetJsonValueExtractor();
+        }
+
         const auto& typeInfo = NYql::NUdf::GetDataTypeInfo(dataSlot);
         return [dataSlot, optional, &typeInfo](simdjson::builtin::ondemand::value jsonValue, NYql::NUdf::TUnboxedValue& resultValue) {
             switch (jsonValue.type()) {
@@ -199,16 +203,7 @@ class TColumnParser {
 
                 case simdjson::builtin::ondemand::json_type::array:
                 case simdjson::builtin::ondemand::json_type::object: {
-                    const auto rawJson = jsonValue.raw_json().value();
-                    if (Y_UNLIKELY(dataSlot != NYql::NUdf::EDataSlot::Json)) {
-                        throw yexception() << "found unexpected nested value (raw: '" << TruncateString(rawJson) << "'), expected data type " <<typeInfo.Name << ", please use Json type for nested values";
-                    }
-                    if (Y_UNLIKELY(!NYql::NDom::IsValidJson(rawJson))) {
-                        throw yexception() << "found bad json value: '" << TruncateString(rawJson) << "'";
-                    }
-                    resultValue = NKikimr::NMiniKQL::MakeString(rawJson);
-                    LockObject(resultValue);
-                    break;
+                    throw yexception() << "found unexpected nested value (raw: '" << TruncateString(jsonValue.raw_json().value()) << "'), expected data type " <<typeInfo.Name << ", please use Json type for nested values";
                 }
 
                 case simdjson::builtin::ondemand::json_type::boolean: {
@@ -230,6 +225,17 @@ class TColumnParser {
         };
     }
 
+    static TParser GetJsonValueExtractor() {
+        return [](simdjson::builtin::ondemand::value jsonValue, NYql::NUdf::TUnboxedValue& resultValue) {
+            const auto rawJson = jsonValue.raw_json().value();
+            if (Y_UNLIKELY(!NYql::NDom::IsValidJson(rawJson))) {
+                throw yexception() << "found bad json value: '" << TruncateString(rawJson) << "'";
+            }
+            resultValue = NKikimr::NMiniKQL::MakeString(rawJson);
+            LockObject(resultValue);
+        };
+    }
+
     template <typename TResult, typename TJsonNumber>
     static NYql::NUdf::TUnboxedValuePod ParseJsonNumber(TJsonNumber number) {
         if (number < std::numeric_limits<TResult>::min() || std::numeric_limits<TResult>::max() < number) {
@@ -239,8 +245,13 @@ class TColumnParser {
     }
 
     static void LockObject(NYql::NUdf::TUnboxedValue& value) {
+        // All UnboxedValue's with type Boxed or String should be locked
+        // because after parsing they will be used under another MKQL allocator in purecalc filters
+
         const i32 numberRefs = value.LockRef();
-        Y_ENSURE(numberRefs == -1 || numberRefs == 1);
+
+        // -1 - value is embbeded or empty, otherwise value should have exactly one ref
+        Y_ENSURE(numberRefs == -1 || numberRefs == 1);  
     }
 
     static TString TruncateString(std::string_view rawString, size_t maxSize = 1_KB) {

ydb/core/fq/libs/row_dispatcher/ut/json_parser_ut.cpp

@@ -185,19 +185,27 @@ Y_UNIT_TEST_SUITE(TJsonParserTests) {
 
         Parser->AddMessages({
             GetMessage(42, R"({"a1": "hello1", "nested": {"key": "value"}})"),
-            GetMessage(43, R"({"a1": "hello2", "nested": ["key1", "key2"]})")
+            GetMessage(43, R"({"a1": "hello2", "nested": ["key1", "key2"]})"),
+            GetMessage(43, R"({"a1": "hello3", "nested": "some string"})"),
+            GetMessage(43, R"({"a1": "hello4", "nested": 123456})")
         });
 
         const auto& result = Parser->Parse();
         ResultNumberValues = result.front().size();
-        UNIT_ASSERT_VALUES_EQUAL(2, ResultNumberValues);
+        UNIT_ASSERT_VALUES_EQUAL(4, ResultNumberValues);
 
         UNIT_ASSERT_VALUES_EQUAL(2, result.size());
         UNIT_ASSERT_VALUES_EQUAL("{\"key\": \"value\"}", TString(result[0][0].AsStringRef()));
         UNIT_ASSERT_VALUES_EQUAL("hello1", TString(result[1][0].AsStringRef()));
 
         UNIT_ASSERT_VALUES_EQUAL("[\"key1\", \"key2\"]", TString(result[0][1].AsStringRef()));
         UNIT_ASSERT_VALUES_EQUAL("hello2", TString(result[1][1].AsStringRef()));
+
+        UNIT_ASSERT_VALUES_EQUAL("\"some string\"", TString(result[0][2].AsStringRef()));
+        UNIT_ASSERT_VALUES_EQUAL("hello3", TString(result[1][2].AsStringRef()));
+
+        UNIT_ASSERT_VALUES_EQUAL("123456", TString(result[0][3].AsStringRef()));
+        UNIT_ASSERT_VALUES_EQUAL("hello4", TString(result[1][3].AsStringRef()));
     }
 
     Y_UNIT_TEST_F(SimpleBooleans, TFixture) {

ydb/tests/fq/yds/test_row_dispatcher.py

@@ -241,7 +241,7 @@ def test_nested_types(self, kikimr, client):
             INSERT INTO {YDS_CONNECTION}.`{self.output_topic}`
             SELECT data FROM {YDS_CONNECTION}.`{self.input_topic}`
                 WITH (format=json_each_row, SCHEMA (time UInt64 NOT NULL, data Json NOT NULL, event String NOT NULL))
-                WHERE event = "event1" or event = "event2";'''
+                WHERE event = "event1" or event = "event2" or event = "event4";'''
 
         query_id = start_yds_query(kikimr, client, sql)
         wait_actor_count(kikimr, "FQ_ROW_DISPATCHER_SESSION", 1)
@@ -251,12 +251,14 @@ def test_nested_types(self, kikimr, client):
             '{"time": 101, "data": {"key": "value", "second_key":"' + large_string + '"}, "event": "event1"}',
             '{"time": 102, "data": ["key1", "key2", "' + large_string + '"], "event": "event2"}',
             '{"time": 103, "data": ["' + large_string + '"], "event": "event3"}',
+            '{"time": 104, "data": "' + large_string + '", "event": "event4"}',
         ]
 
         self.write_stream(data)
         expected = [
             '{"key": "value", "second_key":"' + large_string + '"}',
-            '["key1", "key2", "' + large_string + '"]'
+            '["key1", "key2", "' + large_string + '"]',
+            '"' + large_string + '"'
         ]
         assert self.read_stream(len(expected), topic_path=self.output_topic) == expected