Merge 2ecbeb1 into 2ff8722

ydb/core/formats/arrow/serializer/native.cpp

@@ -113,9 +113,8 @@ NKikimr::TConclusion<std::shared_ptr<arrow::util::Codec>> TNativeSerializer::Bui
     const int levelMin = codec->minimum_compression_level();
     const int levelMax = codec->maximum_compression_level();
     if (levelDef < levelMin || levelMax < levelDef) {
-        return TConclusionStatus::Fail(
-            TStringBuilder() << "incorrect level for codec. have to be: [" << levelMin << ":" << levelMax << "]"
-        );
+        return TConclusionStatus::Fail(TStringBuilder() << "incorrect level for codec `" << arrow::util::Codec::GetCodecAsString(cType)
+                                                        << "`. have to be: [" << levelMin << ":" << levelMax << "]");
     }
     std::shared_ptr<arrow::util::Codec> codecPtr = std::move(NArrow::TStatusValidator::GetValid(arrow::util::Codec::Create(cType, levelDef)));
     return codecPtr;
@@ -182,7 +181,9 @@ NKikimr::TConclusionStatus TNativeSerializer::DoDeserializeFromProto(const NKiki
 void TNativeSerializer::DoSerializeToProto(NKikimrSchemeOp::TOlapColumn::TSerializer& proto) const {
     if (Options.codec) {
         proto.MutableArrowCompression()->SetCodec(NArrow::CompressionToProto(Options.codec->compression_type()));
-        proto.MutableArrowCompression()->SetLevel(Options.codec->compression_level());
+        if (arrow::util::Codec::SupportsCompressionLevel(Options.codec->compression_type())) {
+            proto.MutableArrowCompression()->SetLevel(Options.codec->compression_level());
+        }
     } else {
         proto.MutableArrowCompression()->SetCodec(NArrow::CompressionToProto(arrow::Compression::UNCOMPRESSED));
     }

ydb/core/formats/arrow/serializer/native.h

@@ -103,6 +103,20 @@ class TNativeSerializer: public ISerializer {
         Options.use_threads = false;
         Options.memory_pool = pool;
     }
+
+    arrow::Compression::type GetCodecType() const {
+        if (Options.codec) {
+            return Options.codec->compression_type();
+        }
+        return arrow::Compression::type::UNCOMPRESSED;
+    }
+
+    std::optional<i32> GetCodecLevel() const {
+        if (Options.codec && arrow::util::Codec::SupportsCompressionLevel(Options.codec->compression_type())) {
+            return Options.codec->compression_level();
+        }
+        return {};
+    }
 };
 
 }

ydb/core/formats/arrow/serializer/parsing.cpp

@@ -7,6 +7,18 @@ std::string CompressionToString(const arrow::Compression::type compression) {
     return arrow::util::Codec::GetCodecAsString(compression);
 }
 
+std::string CompressionToString(const NKikimrSchemeOp::EColumnCodec compression) {
+    switch (compression) {
+        case NKikimrSchemeOp::EColumnCodec::ColumnCodecPlain:
+            return "off";
+        case NKikimrSchemeOp::EColumnCodec::ColumnCodecZSTD:
+            return "zstd";
+        case NKikimrSchemeOp::EColumnCodec::ColumnCodecLZ4:
+            return "lz4";
+    }
+    return "";
+}
+
 std::optional<arrow::Compression::type> CompressionFromString(const std::string& compressionStr) {
     auto result = arrow::util::Codec::GetCompressionType(compressionStr);
     if (!result.ok()) {

ydb/core/formats/arrow/serializer/parsing.h

@@ -9,9 +9,9 @@
 namespace NKikimr::NArrow {
 
 std::string CompressionToString(const arrow::Compression::type compression);
+std::string CompressionToString(const NKikimrSchemeOp::EColumnCodec compression);
 std::optional<arrow::Compression::type> CompressionFromString(const std::string& compressionStr);
 
 NKikimrSchemeOp::EColumnCodec CompressionToProto(const arrow::Compression::type compression);
 std::optional<arrow::Compression::type> CompressionFromProto(const NKikimrSchemeOp::EColumnCodec compression);
-
 }

ydb/core/formats/arrow/serializer/utils.cpp

@@ -0,0 +1,29 @@
+#include "parsing.h"
+#include "utils.h"
+
+#include <ydb/library/formats/arrow/validation/validation.h>
+
+#include <contrib/libs/apache/arrow/cpp/src/arrow/util/compression.h>
+
+namespace NKikimr::NArrow {
+bool SupportsCompressionLevel(const arrow::Compression::type compression) {
+    return arrow::util::Codec::SupportsCompressionLevel(compression);
+}
+
+bool SupportsCompressionLevel(const NKikimrSchemeOp::EColumnCodec compression) {
+    return SupportsCompressionLevel(CompressionFromProto(compression).value());
+}
+
+std::optional<int> MinimumCompressionLevel(const arrow::Compression::type compression) {
+    if (!SupportsCompressionLevel(compression)) {
+        return {};
+    }
+    return NArrow::TStatusValidator::GetValid(arrow::util::Codec::MinimumCompressionLevel(compression));
+}
+std::optional<int> MaximumCompressionLevel(const arrow::Compression::type compression) {
+    if (!SupportsCompressionLevel(compression)) {
+        return {};
+    }
+    return NArrow::TStatusValidator::GetValid(arrow::util::Codec::MaximumCompressionLevel(compression));
+}
+}

ydb/core/formats/arrow/serializer/utils.h

@@ -0,0 +1,16 @@
+#pragma once
+
+#include <ydb/core/protos/flat_scheme_op.pb.h>
+
+#include <contrib/libs/apache/arrow/cpp/src/arrow/util/type_fwd.h>
+#include <util/system/yassert.h>
+
+#include <optional>
+
+namespace NKikimr::NArrow {
+bool SupportsCompressionLevel(const arrow::Compression::type compression);
+bool SupportsCompressionLevel(const NKikimrSchemeOp::EColumnCodec compression);
+
+std::optional<int> MinimumCompressionLevel(const arrow::Compression::type compression);
+std::optional<int> MaximumCompressionLevel(const arrow::Compression::type compression);
+}

ydb/core/formats/arrow/serializer/ya.make

@@ -13,6 +13,7 @@ SRCS(
     GLOBAL native.cpp
     stream.cpp
     parsing.cpp
+    utils.cpp
 )
 
 END()

ydb/core/kqp/host/kqp_gateway_proxy.cpp

@@ -132,6 +132,8 @@ bool ConvertCreateTableSettingsToProto(NYql::TKikimrTableMetadataPtr metadata, Y
                 familyProto->set_compression(Ydb::Table::ColumnFamily::COMPRESSION_NONE);
             } else if (to_lower(family.Compression.GetRef()) == "lz4") {
                 familyProto->set_compression(Ydb::Table::ColumnFamily::COMPRESSION_LZ4);
+            } else if (to_lower(family.Compression.GetRef()) == "zstd") {
+                familyProto->set_compression(Ydb::Table::ColumnFamily::COMPRESSION_ZSTD);
             } else {
                 code = Ydb::StatusIds::BAD_REQUEST;
                 error = TStringBuilder() << "Unknown compression '" << family.Compression.GetRef() << "' for a column family";
@@ -383,9 +385,59 @@ bool FillCreateTableDesc(NYql::TKikimrTableMetadataPtr metadata, NKikimrSchemeOp
 }
 
 template <typename T>
-void FillColumnTableSchema(NKikimrSchemeOp::TColumnTableSchema& schema, const T& metadata)
-{
+bool FillColumnTableSchema(NKikimrSchemeOp::TColumnTableSchema& schema, const T& metadata, Ydb::StatusIds::StatusCode& code, TString& error) {
     Y_ENSURE(metadata.ColumnOrder.size() == metadata.Columns.size());
+
+    THashMap<TString, ui32> columnFamiliesByName;
+    ui32 columnFamilyId = 1;
+    for (const auto& family : metadata.ColumnFamilies) {
+        if (family.Data.Defined()) {
+            code = Ydb::StatusIds::BAD_REQUEST;
+            error = TStringBuilder() << "Field `DATA` is not supported for OLAP tables in column family '" << family.Name << "'";
+            return false;
+        }
+        auto columnFamilyIt = columnFamiliesByName.find(family.Name);
+        if (!columnFamilyIt.IsEnd()) {
+            code = Ydb::StatusIds::BAD_REQUEST;
+            error = TStringBuilder() << "Duplicate column family `" << family.Name << '`';
+            return false;
+        }
+        auto familyDescription = schema.AddColumnFamilies();
+        familyDescription->SetName(family.Name);
+        if (familyDescription->GetName() == "default") {
+            familyDescription->SetId(0);
+        } else {
+            familyDescription->SetId(columnFamilyId++);
+        }
+        Y_ENSURE(columnFamiliesByName.emplace(familyDescription->GetName(), familyDescription->GetId()).second);
+        if (family.Compression.Defined()) {
+            NKikimrSchemeOp::EColumnCodec codec;
+            auto codecName = to_lower(family.Compression.GetRef());
+            if (codecName == "off") {
+                codec = NKikimrSchemeOp::EColumnCodec::ColumnCodecPlain;
+            } else if (codecName == "zstd") {
+                codec = NKikimrSchemeOp::EColumnCodec::ColumnCodecZSTD;
+            } else if (codecName == "lz4") {
+                codec = NKikimrSchemeOp::EColumnCodec::ColumnCodecLZ4;
+            } else {
+                code = Ydb::StatusIds::BAD_REQUEST;
+                error = TStringBuilder() << "Unknown compression '" << family.Compression.GetRef() << "' for a column family";
+                return false;
+            }
+            familyDescription->SetColumnCodec(codec);
+        } else {
+            code = Ydb::StatusIds::BAD_REQUEST;
+            error = TStringBuilder() << "Compression is not set for column family'" << family.Name << "'";
+            return false;
+        }
+
+        if (family.CompressionLevel.Defined()) {
+            familyDescription->SetColumnCodecLevel(family.CompressionLevel.GetRef());
+        }
+    }
+
+    schema.SetNextColumnFamilyId(columnFamilyId);
+
     for (const auto& name : metadata.ColumnOrder) {
         auto columnIt = metadata.Columns.find(name);
         Y_ENSURE(columnIt != metadata.Columns.end());
@@ -399,11 +451,29 @@ void FillColumnTableSchema(NKikimrSchemeOp::TColumnTableSchema& schema, const T&
         if (columnType.TypeInfo) {
             *columnDesc.MutableTypeInfo() = *columnType.TypeInfo;
         }
+
+        if (!columnFamiliesByName.empty()) {
+            TString columnFamilyName = "default";
+            ui32 columnFamilyId = 0;
+            if (columnIt->second.Families.size()) {
+                columnFamilyName = *columnIt->second.Families.begin();
+                auto columnFamilyIdIt = columnFamiliesByName.find(columnFamilyName);
+                if (columnFamilyIdIt.IsEnd()) {
+                    code = Ydb::StatusIds::BAD_REQUEST;
+                    error = TStringBuilder() << "Unknown column family `" << columnFamilyName << "` for column `" << columnDesc.GetName() << "`";
+                    return false;
+                }
+                columnFamilyId = columnFamilyIdIt->second;
+            }
+            columnDesc.SetColumnFamilyName(columnFamilyName);
+            columnDesc.SetColumnFamilyId(columnFamilyId);
+        }
     }
 
     for (const auto& keyColumn : metadata.KeyColumnNames) {
         schema.AddKeyColumnNames(keyColumn);
     }
+    return true;
 }
 
 bool FillCreateColumnTableDesc(NYql::TKikimrTableMetadataPtr metadata,
@@ -1661,7 +1731,12 @@ class TKqpGatewayProxy : public IKikimrGateway {
             NKikimrSchemeOp::TColumnTableDescription* tableDesc = schemeTx.MutableCreateColumnTable();
 
             tableDesc->SetName(pathPair.second);
-            FillColumnTableSchema(*tableDesc->MutableSchema(), *metadata);
+            if (!FillColumnTableSchema(*tableDesc->MutableSchema(), *metadata, code, error)) {
+                IKqpGateway::TGenericResult errResult;
+                errResult.AddIssue(NYql::TIssue(error));
+                errResult.SetStatus(NYql::YqlStatusFromYdbStatus(code));
+                return MakeFuture(std::move(errResult));
+            }
 
             if (!FillCreateColumnTableDesc(metadata, *tableDesc, code, error)) {
                 IKqpGateway::TGenericResult errResult;
@@ -1965,7 +2040,22 @@ class TKqpGatewayProxy : public IKikimrGateway {
 
             NKikimrSchemeOp::TColumnTableSchemaPreset* schemaPreset = storeDesc->AddSchemaPresets();
             schemaPreset->SetName("default");
-            FillColumnTableSchema(*schemaPreset->MutableSchema(), settings);
+
+            if (!settings.ColumnFamilies.empty()) {
+                IKqpGateway::TGenericResult errResult;
+                errResult.AddIssue(NYql::TIssue("TableStore does not support column families"));
+                errResult.SetStatus(NYql::YqlStatusFromYdbStatus(Ydb::StatusIds::BAD_REQUEST));
+                return MakeFuture(std::move(errResult));
+            }
+
+            Ydb::StatusIds::StatusCode code;
+            TString error;
+            if (!FillColumnTableSchema(*schemaPreset->MutableSchema(), settings, code, error)) {
+                IKqpGateway::TGenericResult errResult;
+                errResult.AddIssue(NYql::TIssue(error));
+                errResult.SetStatus(NYql::YqlStatusFromYdbStatus(code));
+                return MakeFuture(std::move(errResult));
+            }
 
             if (IsPrepare()) {
                 auto& phyQuery = *SessionCtx->Query().PreparingQuery->MutablePhysicalQuery();

ydb/core/kqp/provider/yql_kikimr_exec.cpp

@@ -182,7 +182,8 @@ namespace {
         return dropGroupSettings;
     }
 
-    TCreateTableStoreSettings ParseCreateTableStoreSettings(TKiCreateTable create, const TTableSettings& settings) {
+    TCreateTableStoreSettings ParseCreateTableStoreSettings(
+        TKiCreateTable create, const TTableSettings& settings, const TVector<TColumnFamily>& columnFamilies) {
         TCreateTableStoreSettings out;
         out.TableStore = TString(create.Table());
         out.ShardsCount = settings.MinPartitions ? *settings.MinPartitions : 0;
@@ -215,6 +216,13 @@ namespace {
                 columnMeta.NotNull = notNull;
             }
 
+            if (columnTuple.Size() > 3) {
+                auto families = columnTuple.Item(3).Cast<TCoAtomList>();
+                for (auto family : families) {
+                    columnMeta.Families.push_back(TString(family.Value()));
+                }
+            }
+
             out.ColumnOrder.push_back(columnName);
             out.Columns.insert(std::make_pair(columnName, columnMeta));
         }
@@ -224,6 +232,7 @@ namespace {
             out.Indexes.push_back(indexDesc);
         }
 #endif
+        out.ColumnFamilies = columnFamilies;
         return out;
     }
 
@@ -1254,8 +1263,8 @@ class TKiSinkCallableExecutionTransformer : public TAsyncCallbackTransformer<TKi
                             TStringBuilder() << "TABLESTORE with not COLUMN store"));
                         return SyncError();
                     }
-                    future = Gateway->CreateTableStore(cluster,
-                        ParseCreateTableStoreSettings(maybeCreate.Cast(), table.Metadata->TableSettings), existingOk);
+                    future = Gateway->CreateTableStore(cluster, ParseCreateTableStoreSettings(maybeCreate.Cast(), table.Metadata->TableSettings,
+                                                                    table.Metadata->ColumnFamilies), existingOk);
                     break;
                 }
                 case ETableType::Table:
@@ -1570,6 +1579,8 @@ class TKiSinkCallableExecutionTransformer : public TAsyncCallbackTransformer<TKi
                                         f->set_compression(Ydb::Table::ColumnFamily::COMPRESSION_NONE);
                                     } else if (to_lower(comp) == "lz4") {
                                         f->set_compression(Ydb::Table::ColumnFamily::COMPRESSION_LZ4);
+                                    } else if (to_lower(comp) == "zstd") {
+                                        f->set_compression(Ydb::Table::ColumnFamily::COMPRESSION_ZSTD);
                                     } else {
                                         auto errText = TStringBuilder() << "Unknown compression '" << comp
                                             << "' for a column family";

ydb/core/kqp/provider/yql_kikimr_gateway.h

@@ -699,6 +699,7 @@ struct TCreateTableStoreSettings {
     TVector<TString> KeyColumnNames;
     TVector<TString> ColumnOrder;
     TVector<TIndexDescription> Indexes;
+    TVector<TColumnFamily> ColumnFamilies;
 };
 
 struct TAlterTableStoreSettings {