YQ-3745 Shared reading: add raw format (#10294)

Author: kardymonds

ydb/library/yql/providers/pq/async_io/dq_pq_rd_read_actor.cpp

@@ -591,7 +591,7 @@ void TDqPqRdReadActor::Handle(NFq::TEvRowDispatcher::TEvMessageBatch::TPtr& ev)
     const NYql::NDqProto::TMessageTransportMeta& meta = ev->Get()->Record.GetTransportMeta();
     SRC_LOG_T("TEvMessageBatch from " << ev->Sender << ", seqNo " << meta.GetSeqNo() << ", ConfirmedSeqNo " << meta.GetConfirmedSeqNo());
     ui64 partitionId = ev->Get()->Record.GetPartitionId();
-    YQL_ENSURE(Sessions.count(partitionId), "Unknown partition id");
+    YQL_ENSURE(Sessions.count(partitionId), "Unknown partition id " << partitionId);
     auto it = Sessions.find(partitionId);
     if (it == Sessions.end()) {
         Stop("Wrong session data");

ydb/library/yql/providers/pq/provider/yql_pq_dq_integration.cpp

@@ -268,11 +268,12 @@ class TPqDqIntegration: public TDqIntegrationBase {
                     }
                 }
 
-                //sharedReading = true;
-                sharedReading = sharedReading && (format == "json_each_row");
+                sharedReading = sharedReading && (format == "json_each_row" || (format == "raw"));
                 TString predicateSql = NYql::FormatWhere(predicateProto);
                 if (sharedReading) {
-                    srcDesc.SetPredicate(predicateSql);
+                    if (format == "json_each_row") {
+                        srcDesc.SetPredicate(predicateSql);
+                    }
                     srcDesc.SetSharedReading(true);
                 }
                 protoSettings.PackFrom(srcDesc);

ydb/tests/fq/yds/test_row_dispatcher.py

@@ -83,29 +83,36 @@ def wait_row_dispatcher_sensor_value(kikimr, sensor, expected_count, exact_match
 class TestPqRowDispatcher(TestYdsBase):
 
     @yq_v1
-    def test_read_raw_format_without_row_dispatcher(self, kikimr, client):
+    def test_read_raw_format_with_row_dispatcher(self, kikimr, client):
         client.create_yds_connection(
             YDS_CONNECTION, os.getenv("YDB_DATABASE"), os.getenv("YDB_ENDPOINT"), shared_reading=True
         )
         self.init_topics("test_read_raw_format_without_row_dispatcher", create_output=False)
-
         output_topic = "pq_test_pq_read_write_output"
-
         create_stream(output_topic, partitions_count=1)
         create_read_rule(output_topic, self.consumer_name)
 
-        sql = Rf'''
-            INSERT INTO {YDS_CONNECTION}.`{output_topic}`
-            SELECT * FROM {YDS_CONNECTION}.`{self.input_topic}`;'''
+        sql1 = Rf'''INSERT INTO {YDS_CONNECTION}.`{output_topic}`
+                    SELECT * FROM {YDS_CONNECTION}.`{self.input_topic}` WITH (format=raw, SCHEMA (data String NOT NULL));'''
 
-        query_id = start_yds_query(kikimr, client, sql)
+        query_id = start_yds_query(kikimr, client, sql1)
         data = ['{"time" = 101;}', '{"time" = 102;}']
 
         self.write_stream(data)
-        expected = data
-        assert self.read_stream(len(expected), topic_path=output_topic) == expected
+        assert self.read_stream(len(data), topic_path=output_topic) == data
+        wait_actor_count(kikimr, "FQ_ROW_DISPATCHER_SESSION", 1)
+        stop_yds_query(client, query_id)
+        
+        sql2 = Rf'''INSERT INTO {YDS_CONNECTION}.`{output_topic}`
+                    SELECT * FROM {YDS_CONNECTION}.`{self.input_topic}` WITH (format=raw, SCHEMA (data String NOT NULL))
+                    WHERE data != "romashka";'''
 
-        wait_actor_count(kikimr, "FQ_ROW_DISPATCHER_SESSION", 0)
+        query_id = start_yds_query(kikimr, client, sql2)
+        data = ['{"time" = 103;}', '{"time" = 104;}']
+
+        self.write_stream(data)
+        assert self.read_stream(len(data), topic_path=output_topic) == data
+        wait_actor_count(kikimr, "FQ_ROW_DISPATCHER_SESSION", 1)
         stop_yds_query(client, query_id)
 
     @yq_v1
@@ -193,7 +200,7 @@ def test_scheme_error(self, kikimr, client):
 
         client.wait_query_status(query_id, fq.QueryMeta.FAILED)
         issues = str(client.describe_query(query_id).result.query.issue)
-        assert "Failed to unwrap empty optional" in issues, "Incorrect Issues: " + issues
+        assert "Cannot parse JSON string" in issues, "Incorrect Issues: " + issues
 
         wait_actor_count(kikimr, "DQ_PQ_READ_ACTOR", 0)
         wait_actor_count(kikimr, "FQ_ROW_DISPATCHER_SESSION", 0)