fixed issues

Author: evanevanevanevannnn

ydb/core/external_sources/object_storage/inference/arrow_fetcher.cpp

@@ -152,30 +152,40 @@ class TArrowFileFetcher : public NActors::TActorBootstrapped<TArrowFileFetcher>
     // Cutting file
 
     TMaybe<TString> DecompressFile(const TString& data, const TRequest& request, const NActors::TActorContext& ctx) {
-        auto dataBuffer = NDB::ReadBufferFromString(data);
-        auto decompressorBuffer = NYql::MakeDecompressor(dataBuffer, *DecompressionFormat_);
-        if (!decompressorBuffer) {
-            auto error = MakeError(
+        try {
+            auto dataBuffer = NDB::ReadBufferFromString(data);
+            auto decompressorBuffer = NYql::MakeDecompressor(dataBuffer, *DecompressionFormat_);
+            if (!decompressorBuffer) {
+                auto error = MakeError(
+                    request.Path,
+                    NFq::TIssuesIds::INTERNAL_ERROR,
+                    TStringBuilder{} << "unknown compression: " << *DecompressionFormat_ << ". Use one of: gzip, zstd, lz4, brotli, bzip2, xz" 
+                );
+                SendError(ctx, error);
+                return {};
+            }
+
+            TStringBuilder decompressedData;
+            while (!decompressorBuffer->eof() && decompressedData.size() < 10_MB) {
+                decompressorBuffer->nextIfAtEnd();
+                size_t maxDecompressedChunkSize = std::min(
+                    decompressorBuffer->available(),                
+                    10_MB - decompressedData.size()
+                );
+                TString decompressedChunk{maxDecompressedChunkSize, ' '};
+                decompressorBuffer->read(&decompressedChunk.front(), maxDecompressedChunkSize);
+                decompressedData << decompressedChunk;
+            }
+            return std::move(decompressedData);
+        } catch (const yexception& error) {
+            auto errorEv = MakeError(
                 request.Path,
                 NFq::TIssuesIds::INTERNAL_ERROR,
-                TStringBuilder{} << "invalid decompression format: " << *DecompressionFormat_
+                TStringBuilder{} << "couldn't decompress file, check format and compression params: " << error.AsStrBuf()
             );
-            SendError(ctx, error);
+            SendError(ctx, errorEv);
             return {};
         }
-
-        TString decompressedData;
-        while (!decompressorBuffer->eof() && decompressedData.size() < 10_MB) {
-            decompressorBuffer->nextIfAtEnd();
-            size_t maxDecompressedChunkSize = std::min(
-                decompressorBuffer->available(),                
-                10_MB - decompressedData.size()
-            );
-            TString decompressedChunk{maxDecompressedChunkSize, ' '};
-            decompressorBuffer->read(&decompressedChunk.front(), maxDecompressedChunkSize);
-            decompressedData += decompressedChunk;
-        }
-        return std::move(decompressedData);
     }
 
     std::shared_ptr<arrow::io::RandomAccessFile> CleanupCsvFile(const TString& data, const TRequest& request, const arrow::csv::ParseOptions& options, const NActors::TActorContext& ctx) {

ydb/core/external_sources/object_storage/inference/ut/arrow_inference_ut.cpp

@@ -50,7 +50,7 @@ class ArrowInferenceTest : public testing::Test {
 
     NActors::TActorId RegisterInferencinator(TStringBuf formatStr) {
         auto format = NInference::ConvertFileFormat(formatStr);
-        auto arrowFetcher = ActorSystem.Register(NInference::CreateArrowFetchingActor(S3ActorId, format), 1);
+        auto arrowFetcher = ActorSystem.Register(NInference::CreateArrowFetchingActor(S3ActorId, format, {}), 1);
         return ActorSystem.Register(NInference::CreateArrowInferencinator(arrowFetcher, format, {}), 1);
     }
 

ydb/core/external_sources/object_storage/inference/ya.make

@@ -15,6 +15,7 @@ PEERDIR(
 
     ydb/core/external_sources/object_storage
 
+    ydb/library/yql/providers/s3/compressors
     ydb/library/yql/udfs/common/clickhouse/client
 )
 

ydb/tests/fq/s3/test_compressions.py

@@ -204,3 +204,39 @@ def test_invalid_compression(self, kikimr, s3, client, unique_prefix):
         assert (
             "Unknown compression: some_compression. Use one of: gzip, zstd, lz4, brotli, bzip2, xz" in describe_string
         )
+    
+    @yq_v2
+    @pytest.mark.parametrize("client", [{"folder_id": "my_folder"}], indirect=True)
+    def test_invalid_compression_inference(self, kikimr, s3, client, unique_prefix):
+        resource = boto3.resource(
+            "s3", endpoint_url=s3.s3_url, aws_access_key_id="key", aws_secret_access_key="secret_key"
+        )
+
+        bucket = resource.Bucket("fbucket")
+        bucket.create(ACL='public-read')
+
+        s3_client = boto3.client(
+            "s3", endpoint_url=s3.s3_url, aws_access_key_id="key", aws_secret_access_key="secret_key"
+        )
+
+        fruits = R'''[{"name" : "banana", "price" : 3, "weight" : 100}]'''
+        s3_client.put_object(Body=fruits, Bucket='fbucket', Key='fruits.json', ContentType='text/plain')
+        kikimr.control_plane.wait_bootstrap(1)
+
+        storage_connection_name = unique_prefix + "fruitbucket"
+        client.create_storage_connection(storage_connection_name, "fbucket")
+
+        sql = fR'''
+            SELECT *
+            FROM `{storage_connection_name}`.`fruits.json`
+            WITH (format=csv_with_names, compression="gzip", with_infer="true");
+            '''
+
+        query_id = client.create_query("simple", sql, type=fq.QueryContent.QueryType.ANALYTICS).result.query_id
+        client.wait_query_status(query_id, fq.QueryMeta.FAILED)
+        describe_result = client.describe_query(query_id).result
+        logging.debug("Describe result: {}".format(describe_result))
+        describe_string = "{}".format(describe_result)
+        assert (
+            "couldn\\'t decompress file, check format and compression params:" in describe_string
+        )