Decompression support for schema inference (#7632)

Author: evanevanevanevannnn

ydb/core/external_sources/object_storage.cpp

@@ -349,7 +349,7 @@ struct TObjectStorageExternalSource : public IExternalSource {
         meta->Attributes.erase("withinfer");
 
         auto fileFormat = NObjectStorage::NInference::ConvertFileFormat(*format);
-        auto arrowFetcherId = ActorSystem->Register(NObjectStorage::NInference::CreateArrowFetchingActor(s3FetcherId, fileFormat));
+        auto arrowFetcherId = ActorSystem->Register(NObjectStorage::NInference::CreateArrowFetchingActor(s3FetcherId, fileFormat, meta->Attributes));
         auto arrowInferencinatorId = ActorSystem->Register(NObjectStorage::NInference::CreateArrowInferencinator(arrowFetcherId, fileFormat, meta->Attributes));
 
         return afterListing.Apply([arrowInferencinatorId, meta, actorSystem = ActorSystem](const NThreading::TFuture<TString>& pathFut) {

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

@@ -13,17 +13,24 @@
 #include <ydb/core/external_sources/object_storage/events.h>
 #include <ydb/library/actors/core/actor_bootstrapped.h>
 #include <ydb/library/actors/core/hfunc.h>
+#include <ydb/library/yql/providers/s3/compressors/factory.h>
+#include <ydb/library/yql/udfs/common/clickhouse/client/src/IO/ReadBufferFromString.h>
 
 namespace NKikimr::NExternalSource::NObjectStorage::NInference {
 
 class TArrowFileFetcher : public NActors::TActorBootstrapped<TArrowFileFetcher> {
     static constexpr uint64_t PrefixSize = 10_MB;
 public:
-    TArrowFileFetcher(NActors::TActorId s3FetcherId, EFileFormat format)
+    TArrowFileFetcher(NActors::TActorId s3FetcherId, EFileFormat format, const THashMap<TString, TString>& params)
         : S3FetcherId_{s3FetcherId}
         , Format_{format}
     {
         Y_ABORT_UNLESS(IsArrowInferredFormat(Format_));
+        
+        auto decompression = params.FindPtr("compression");
+        if (decompression) {
+            DecompressionFormat_ = *decompression;
+        }
     }
 
     void Bootstrap() {
@@ -67,6 +74,15 @@ class TArrowFileFetcher : public NActors::TActorBootstrapped<TArrowFileFetcher>
 
         const auto& request = requestIt->second;
 
+        TString data = std::move(response.Data);
+        if (DecompressionFormat_) {
+            auto decompressedData = DecompressFile(data, request, ctx);
+            if (!decompressedData) {
+                return;
+            }
+            data = std::move(*decompressedData);
+        }
+
         std::shared_ptr<arrow::io::RandomAccessFile> file;
         switch (Format_) {
             case EFileFormat::CsvWithNames:
@@ -76,7 +92,7 @@ class TArrowFileFetcher : public NActors::TActorBootstrapped<TArrowFileFetcher>
                 if (Format_ == EFileFormat::TsvWithNames) {
                     options.delimiter = '\t';
                 }
-                file = CleanupCsvFile(response.Data, request, options, ctx);
+                file = CleanupCsvFile(data, request, options, ctx);
                 ctx.Send(request.Requester, new TEvArrowFile(std::move(file), request.Path));
                 break;
             }
@@ -135,6 +151,43 @@ class TArrowFileFetcher : public NActors::TActorBootstrapped<TArrowFileFetcher>
 
     // Cutting file
 
+    TMaybe<TString> DecompressFile(const TString& data, const TRequest& request, const NActors::TActorContext& ctx) {
+        try {
+            NDB::ReadBufferFromString dataBuffer(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{} << "couldn't decompress file, check compression params: " << error.what()
+            );
+            SendError(ctx, errorEv);
+            return {};
+        }
+    }
+
     std::shared_ptr<arrow::io::RandomAccessFile> CleanupCsvFile(const TString& data, const TRequest& request, const arrow::csv::ParseOptions& options, const NActors::TActorContext& ctx) {
         auto chunker = arrow::csv::MakeChunker(options);
         std::shared_ptr<arrow::Buffer> whole, partial;
@@ -183,10 +236,11 @@ class TArrowFileFetcher : public NActors::TActorBootstrapped<TArrowFileFetcher>
     // Fields
     NActors::TActorId S3FetcherId_;
     EFileFormat Format_;
+    TMaybe<TString> DecompressionFormat_;
     std::unordered_map<TString, TRequest> InflightRequests_; // Path -> Request
 };
 
-NActors::IActor* CreateArrowFetchingActor(NActors::TActorId s3FetcherId, EFileFormat format) {
-    return new TArrowFileFetcher{s3FetcherId, format};
+NActors::IActor* CreateArrowFetchingActor(NActors::TActorId s3FetcherId, EFileFormat format, const THashMap<TString, TString>& params) {
+    return new TArrowFileFetcher{s3FetcherId, format, params};
 }
 } // namespace NKikimr::NExternalSource::NObjectStorage::NInference

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

@@ -5,5 +5,5 @@
 
 namespace NKikimr::NExternalSource::NObjectStorage::NInference {
 
-NActors::IActor* CreateArrowFetchingActor(NActors::TActorId s3FetcherId, EFileFormat format);
+NActors::IActor* CreateArrowFetchingActor(NActors::TActorId s3FetcherId, EFileFormat format, const THashMap<TString, TString>& params);
 } // namespace NKikimr::NExternalSource::NObjectStorage::NInference

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/ut/ya.make

@@ -1,6 +1,7 @@
 GTEST()
 
 PEERDIR(
+    ydb/library/yql/public/udf/service/stub
     ydb/core/external_sources/object_storage/inference
     ydb/core/external_sources/object_storage
     ydb/core/tx/scheme_board

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

@@ -1,5 +1,11 @@
 LIBRARY()
 
+ADDINCL(
+    ydb/library/yql/udfs/common/clickhouse/client/base
+    ydb/library/yql/udfs/common/clickhouse/client/base/pcg-random
+    ydb/library/yql/udfs/common/clickhouse/client/src
+)
+
 SRCS(
     arrow_fetcher.cpp
     arrow_inferencinator.cpp
@@ -9,6 +15,8 @@ PEERDIR(
     contrib/libs/apache/arrow
 
     ydb/core/external_sources/object_storage
+
+    ydb/library/yql/providers/s3/compressors
 )
 
 END()

ydb/tests/fq/s3/test_compression_data/test.csv.br

@@ -10,7 +10,7 @@
 import ydb.public.api.protos.draft.fq_pb2 as fq
 
 import ydb.tests.fq.s3.s3_helpers as s3_helpers
-from ydb.tests.tools.fq_runner.kikimr_utils import yq_all
+from ydb.tests.tools.fq_runner.kikimr_utils import yq_all, yq_v2
 
 
 class TestS3Compressions:
@@ -33,6 +33,26 @@ def validate_result(self, result_set):
         assert result_set.rows[0].items[0].bytes_value == b"yq"
         assert result_set.rows[0].items[1].int32_value == 0
         assert result_set.rows[0].items[2].bytes_value == b"abc"
+    
+    def validate_result_csv(self, result_set):
+        logging.debug(str(result_set))
+        assert len(result_set.columns) == 3
+        assert result_set.columns[0].name == "a"
+        assert result_set.columns[0].type.optional_type.item.type_id == ydb.Type.INT64
+        assert result_set.columns[1].name == "b"
+        assert result_set.columns[1].type.type_id == ydb.Type.UTF8
+        assert result_set.columns[2].name == "c"
+        assert result_set.columns[2].type.optional_type.item.type_id == ydb.Type.DOUBLE
+        assert len(result_set.rows) == 3
+        assert result_set.rows[0].items[0].int64_value == 1
+        assert result_set.rows[0].items[1].text_value == "hello"
+        assert result_set.rows[0].items[2].double_value == 0.5
+        assert result_set.rows[1].items[0].int64_value == 2
+        assert result_set.rows[1].items[1].text_value == "world"
+        assert result_set.rows[1].items[2].double_value == 0.25
+        assert result_set.rows[2].items[0].int64_value == 3
+        assert result_set.rows[2].items[1].text_value == "!"
+        assert result_set.rows[2].items[2].double_value == 0.125
 
     @yq_all
     @pytest.mark.parametrize(
@@ -70,6 +90,38 @@ def test_compression(self, kikimr, s3, client, filename, compression, unique_pre
         result_set = data.result.result_set
         self.validate_result(result_set)
 
+    @yq_v2
+    @pytest.mark.parametrize(
+        "filename, compression",
+        [
+            ("test.csv.gz", "gzip"),
+            ("test.csv.lz4", "lz4"),
+            ("test.csv.br", "brotli"),
+            ("test.csv.bz2", "bzip2"),
+            ("test.csv.zst", "zstd"),
+            ("test.csv.xz", "xz"),
+        ],
+    )
+    def test_compression_inference(self, kikimr, s3, client, filename, compression, unique_prefix):
+        self.create_bucket_and_upload_file(filename, s3, kikimr)
+        storage_connection_name = unique_prefix + "fruitbucket"
+        client.create_storage_connection(storage_connection_name, "fbucket")
+        
+        sql = '''
+            SELECT *
+            FROM `{}`.`{}`
+            WITH (format=csv_with_names, compression="{}", with_infer="true");
+            '''.format(
+            storage_connection_name, filename, compression
+        )
+
+        query_id = client.create_query("simple", sql, type=fq.QueryContent.QueryType.ANALYTICS).result.query_id
+        client.wait_query_status(query_id, fq.QueryMeta.COMPLETED)
+
+        data = client.get_result_data(query_id)
+        result_set = data.result.result_set
+        self.validate_result_csv(result_set)
+
     @yq_all
     @pytest.mark.parametrize(
         "filename, compression",
@@ -152,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 compression params:" in describe_string
+        )