Optimize alter-table verb

commit_hash:dbdf35744550a5a2b4a9b94f9b325b60f1dde718

4 files changed, 218 insertions, 26 deletions

diff --git a/yt/yt/client/table_client/logical_type.cpp b/yt/yt/client/table_client/logical_type.cpp
index 5d9a5f75b6..46ae7221da 100644
--- a/yt/yt/client/table_client/logical_type.cpp
+++ b/yt/yt/client/table_client/logical_type.cpp
@@ -320,11 +320,18 @@ TDecimalLogicalType::TDecimalLogicalType(int precision, int scale)
     , Scale_(scale)
 { }
 
-size_t TDecimalLogicalType::GetMemoryUsage() const
+i64 TDecimalLogicalType::GetMemoryUsage() const
 {
     return sizeof(*this);
 }
 
+i64 TDecimalLogicalType::GetMemoryUsage(i64 limit) const
+{
+    YT_ASSERT(limit > 0);
+
+    return sizeof(*this);
+}
+
 int TDecimalLogicalType::GetTypeComplexity() const
 {
     return 1;
@@ -372,7 +379,7 @@ std::optional<ESimpleLogicalValueType> TOptionalLogicalType::Simplify() const
     }
 }
 
-size_t TOptionalLogicalType::GetMemoryUsage() const
+i64 TOptionalLogicalType::GetMemoryUsage() const
 {
     if (Element_->GetMetatype() == ELogicalMetatype::Simple) {
         // All optionals of simple logical types are singletons and therefore we assume they use no space.
@@ -382,6 +389,20 @@ size_t TOptionalLogicalType::GetMemoryUsage() const
     }
 }
 
+i64 TOptionalLogicalType::GetMemoryUsage(i64 limit) const
+{
+    YT_ASSERT(limit > 0);
+
+    if (Element_->GetMetatype() == ELogicalMetatype::Simple) {
+        // NB: see TOptionalLogicalType::GetMemoryUsage().
+        return 0;
+    } else if (auto sizeOfThis = static_cast<i64>(sizeof(*this)); sizeOfThis >= limit) {
+        return sizeof(*this);
+    } else {
+        return sizeOfThis + Element_->GetMemoryUsage(limit - sizeOfThis);
+    }
+}
+
 int TOptionalLogicalType::GetTypeComplexity() const
 {
     if (Element_->GetMetatype() == ELogicalMetatype::Simple) {
@@ -406,12 +427,19 @@ TSimpleLogicalType::TSimpleLogicalType(ESimpleLogicalValueType element)
     , Element_(element)
 { }
 
-size_t TSimpleLogicalType::GetMemoryUsage() const
+i64 TSimpleLogicalType::GetMemoryUsage() const
 {
     // All simple logical types are singletons and therefore we assume they use no space.
     return 0;
 }
 
+i64 TSimpleLogicalType::GetMemoryUsage(i64 limit) const
+{
+    YT_ASSERT(limit > 0);
+
+    return 0;
+}
+
 int TSimpleLogicalType::GetTypeComplexity() const
 {
     return 1;
@@ -439,11 +467,22 @@ TListLogicalType::TListLogicalType(TLogicalTypePtr element)
     , Element_(std::move(element))
 { }
 
-size_t TListLogicalType::GetMemoryUsage() const
+i64 TListLogicalType::GetMemoryUsage() const
 {
     return sizeof(*this) + Element_->GetMemoryUsage();
 }
 
+i64 TListLogicalType::GetMemoryUsage(i64 limit) const
+{
+    YT_ASSERT(limit > 0);
+
+    if (auto sizeOfThis = static_cast<i64>(sizeof(*this)); sizeOfThis >= limit) {
+        return sizeOfThis;
+    } else {
+        return sizeOfThis + Element_->GetMemoryUsage(limit - sizeOfThis);
+    }
+}
+
 int TListLogicalType::GetTypeComplexity() const
 {
     return 1 + Element_->GetTypeComplexity();
@@ -678,14 +717,28 @@ TStructLogicalTypeBase::TStructLogicalTypeBase(ELogicalMetatype metatype, std::v
     , Fields_(std::move(fields))
 { }
 
-size_t TStructLogicalTypeBase::GetMemoryUsage() const
+i64 TStructLogicalTypeBase::GetMemoryUsage() const
 {
-    size_t result = sizeof(*this);
-    result += sizeof(TStructField) * Fields_.size();
+    auto usage = static_cast<i64>(sizeof(*this));
+    usage += sizeof(TStructField) * Fields_.size();
     for (const auto& field : Fields_) {
-        result += field.Type->GetMemoryUsage();
+        usage += field.Type->GetMemoryUsage();
     }
-    return result;
+    return usage;
+}
+
+i64 TStructLogicalTypeBase::GetMemoryUsage(i64 limit) const
+{
+    YT_ASSERT(limit > 0);
+
+    auto usage = static_cast<i64>(sizeof(*this) + sizeof(TStructField) * Fields_.size());
+    for (const auto& field : Fields_) {
+        if (usage >= limit) {
+            return usage;
+        }
+        usage += field.Type->GetMemoryUsage(limit - usage);
+    }
+    return usage;
 }
 
 int TStructLogicalTypeBase::GetTypeComplexity() const
@@ -736,14 +789,28 @@ TTupleLogicalTypeBase::TTupleLogicalTypeBase(ELogicalMetatype metatype, std::vec
     , Elements_(std::move(elements))
 { }
 
-size_t TTupleLogicalTypeBase::GetMemoryUsage() const
+i64 TTupleLogicalTypeBase::GetMemoryUsage() const
 {
-    size_t result = sizeof(*this);
-    result += sizeof(TLogicalTypePtr) * Elements_.size();
+    auto usage = static_cast<i64>(sizeof(*this));
+    usage += sizeof(TLogicalTypePtr) * Elements_.size();
     for (const auto& element : Elements_) {
-        result += element->GetMemoryUsage();
+        usage += element->GetMemoryUsage();
     }
-    return result;
+    return usage;
+}
+
+i64 TTupleLogicalTypeBase::GetMemoryUsage(i64 limit) const
+{
+    YT_ASSERT(limit > 0);
+
+    auto usage = static_cast<i64>(sizeof(*this) + sizeof(TLogicalTypePtr) * Elements_.size());
+    for (const auto& element : Elements_) {
+        if (usage >= limit) {
+            return usage;
+        }
+        usage += element->GetMemoryUsage(limit - usage);
+    }
+    return usage;
 }
 
 int TTupleLogicalTypeBase::GetTypeComplexity() const
@@ -795,11 +862,27 @@ TDictLogicalType::TDictLogicalType(TLogicalTypePtr key, TLogicalTypePtr value)
     , Value_(std::move(value))
 { }
 
-size_t TDictLogicalType::GetMemoryUsage() const
+i64 TDictLogicalType::GetMemoryUsage() const
 {
     return sizeof(*this) + Key_->GetMemoryUsage() + Value_->GetMemoryUsage();
 }
 
+i64 TDictLogicalType::GetMemoryUsage(i64 limit) const
+{
+    YT_ASSERT(limit > 0);
+
+    auto usage = static_cast<i64>(sizeof(*this));
+    if (usage >= limit) {
+        return usage;
+    }
+    usage += Key_->GetMemoryUsage(limit - usage);
+    if (usage >= limit) {
+        return usage;
+    }
+    usage += Value_->GetMemoryUsage(limit - usage);
+    return usage;
+}
+
 int TDictLogicalType::GetTypeComplexity() const
 {
     return 1 + Key_->GetTypeComplexity() + Value_->GetTypeComplexity();
@@ -877,11 +960,23 @@ TTaggedLogicalType::TTaggedLogicalType(TString tag, NYT::NTableClient::TLogicalT
     , Element_(std::move(element))
 { }
 
-size_t TTaggedLogicalType::GetMemoryUsage() const
+i64 TTaggedLogicalType::GetMemoryUsage() const
 {
     return sizeof(*this) + GetElement()->GetMemoryUsage();
 }
 
+i64 TTaggedLogicalType::GetMemoryUsage(i64 limit) const
+{
+    YT_ASSERT(limit > 0);
+
+    auto usage = static_cast<i64>(sizeof(*this));
+    if (usage >= limit) {
+        return usage;
+    }
+    usage += GetElement()->GetMemoryUsage(limit - usage);
+    return usage;
+}
+
 int TTaggedLogicalType::GetTypeComplexity() const
 {
     return 1 + GetElement()->GetTypeComplexity();
diff --git a/yt/yt/client/table_client/logical_type.h b/yt/yt/client/table_client/logical_type.h
index 1a0b2dd0b5..2c9a6afbfa 100644
--- a/yt/yt/client/table_client/logical_type.h
+++ b/yt/yt/client/table_client/logical_type.h
@@ -71,7 +71,8 @@ public:
     const TTaggedLogicalType& AsTaggedTypeRef() const;
     Y_FORCE_INLINE const TTaggedLogicalType& UncheckedAsTaggedTypeRef() const;
 
-    virtual size_t GetMemoryUsage() const = 0;
+    virtual i64 GetMemoryUsage() const = 0;
+    virtual i64 GetMemoryUsage(i64 limit) const = 0;
     virtual int GetTypeComplexity() const = 0;
 
     // This function doesn't validate children of current node.
@@ -175,7 +176,8 @@ public:
 public:
     TDecimalLogicalType(int precision, int scale);
 
-    size_t GetMemoryUsage() const override;
+    i64 GetMemoryUsage() const override;
+    i64 GetMemoryUsage(i64 limit) const override;
     int GetTypeComplexity() const override;
     void ValidateNode(const TWalkContext& context) const override;
     bool IsNullable() const override;
@@ -203,7 +205,8 @@ public:
     // Cached value of GetElement()->IsNullable(), useful for performance reasons.
     Y_FORCE_INLINE bool IsElementNullable() const;
 
-    size_t GetMemoryUsage() const override;
+    i64 GetMemoryUsage() const override;
+    i64 GetMemoryUsage(i64 limit) const override;
     int GetTypeComplexity() const override;
     void ValidateNode(const TWalkContext& context) const override;
     bool IsNullable() const override;
@@ -223,7 +226,8 @@ public:
 
     Y_FORCE_INLINE ESimpleLogicalValueType GetElement() const;
 
-    size_t GetMemoryUsage() const override;
+    i64 GetMemoryUsage() const override;
+    i64 GetMemoryUsage(i64 limit) const override;
     int GetTypeComplexity() const override;
     void ValidateNode(const TWalkContext& context) const override;
     bool IsNullable() const override;
@@ -242,7 +246,8 @@ public:
 
     Y_FORCE_INLINE const TLogicalTypePtr& GetElement() const;
 
-    size_t GetMemoryUsage() const override;
+    i64 GetMemoryUsage() const override;
+    i64 GetMemoryUsage(i64 limit) const override;
     int GetTypeComplexity() const override;
     void ValidateNode(const TWalkContext& context) const override;
     bool IsNullable() const override;
@@ -302,7 +307,8 @@ public:
     TStructLogicalTypeBase(ELogicalMetatype metatype, std::vector<TStructField> fields);
     Y_FORCE_INLINE const std::vector<TStructField>& GetFields() const;
 
-    size_t GetMemoryUsage() const override;
+    i64 GetMemoryUsage() const override;
+    i64 GetMemoryUsage(i64 limit) const override;
     int GetTypeComplexity() const override;
     void ValidateNode(const TWalkContext& context) const override;
     bool IsNullable() const override;
@@ -321,7 +327,8 @@ public:
 
     Y_FORCE_INLINE const std::vector<TLogicalTypePtr>& GetElements() const;
 
-    size_t GetMemoryUsage() const override;
+    i64 GetMemoryUsage() const override;
+    i64 GetMemoryUsage(i64 limit) const override;
     int GetTypeComplexity() const override;
     void ValidateNode(const TWalkContext& context) const override;
     bool IsNullable() const override;
@@ -377,7 +384,8 @@ public:
     Y_FORCE_INLINE const TLogicalTypePtr& GetKey() const;
     Y_FORCE_INLINE const TLogicalTypePtr& GetValue() const;
 
-    size_t GetMemoryUsage() const override;
+    i64 GetMemoryUsage() const override;
+    i64 GetMemoryUsage(i64 limit) const override;
     int GetTypeComplexity() const override;
     void ValidateNode(const TWalkContext& context) const override;
     bool IsNullable() const override;
@@ -398,7 +406,8 @@ public:
     Y_FORCE_INLINE const TString& GetTag() const;
     Y_FORCE_INLINE const TLogicalTypePtr& GetElement() const;
 
-    size_t GetMemoryUsage() const override;
+    i64 GetMemoryUsage() const override;
+    i64 GetMemoryUsage(i64 limit) const override;
     int GetTypeComplexity() const override;
     void ValidateNode(const TWalkContext& context) const override;
     bool IsNullable() const override;
diff --git a/yt/yt/client/table_client/schema.cpp b/yt/yt/client/table_client/schema.cpp
index dda743a92d..7416c7462e 100644
--- a/yt/yt/client/table_client/schema.cpp
+++ b/yt/yt/client/table_client/schema.cpp
@@ -16,6 +16,8 @@
 #include <yt/yt/core/ytree/fluent.h>
 #include <yt/yt/core/ytree/yson_struct.h>
 
+#include <yt/yt/core/actions/new_with_offloaded_dtor.h>
+
 #include <yt/yt_proto/yt/client/table_chunk_format/proto/chunk_meta.pb.h>
 #include <yt/yt_proto/yt/client/table_chunk_format/proto/wire_protocol.pb.h>
 
@@ -280,6 +282,26 @@ i64 TColumnSchema::GetMemoryUsage() const
         (Group_ ? Group_->size() : 0);
 }
 
+i64 TColumnSchema::GetMemoryUsage(i64 limit) const
+{
+    YT_ASSERT(limit > 0);
+
+    auto usage = static_cast<i64>(
+        sizeof(TColumnSchema) +
+        StableName_.Underlying().size() +
+        Name_.size() +
+        (Lock_ ? Lock_->size() : 0) +
+        (Expression_ ? Expression_->size() : 0) +
+        (Aggregate_ ? Aggregate_->size() : 0) +
+        (Group_ ? Group_->size() : 0));
+
+    if (usage >= limit) {
+        return usage;
+    }
+
+    return usage + LogicalType_->GetMemoryUsage(limit - usage);
+}
+
 bool TColumnSchema::IsOfV1Type() const
 {
     return IsOfV1Type_;
@@ -1430,13 +1452,28 @@ void TTableSchema::Load(TStreamLoadContext& context)
 
 i64 TTableSchema::GetMemoryUsage() const
 {
-    i64 usage = sizeof(TTableSchema);
+    auto usage = static_cast<i64>(sizeof(TTableSchema));
     for (const auto& column : Columns()) {
         usage += column.GetMemoryUsage();
     }
     return usage;
 }
 
+i64 TTableSchema::GetMemoryUsage(i64 limit) const
+{
+    YT_ASSERT(limit > 0);
+
+    auto usage = static_cast<i64>(sizeof(TTableSchema));
+    for (const auto& column : Columns()) {
+        if (usage >= limit) {
+            return usage;
+        }
+
+        usage += column.GetMemoryUsage(limit - usage);
+    }
+    return usage;
+}
+
 TKeyColumnTypes TTableSchema::GetKeyColumnTypes() const
 {
     TKeyColumnTypes result(KeyColumnCount_);
@@ -1574,6 +1611,35 @@ void PrintTo(const TTableSchema& tableSchema, std::ostream* os)
 
 ////////////////////////////////////////////////////////////////////////////////
 
+TTableSchemaTruncatedFormatter::TTableSchemaTruncatedFormatter(
+    const TTableSchemaPtr& schema,
+    i64 memoryLimit)
+    : Schema_(schema.Get())
+    , Limit_(memoryLimit)
+{
+    YT_ASSERT(Limit_ >= 0);
+}
+
+void TTableSchemaTruncatedFormatter::operator()(TStringBuilderBase* builder) const
+{
+    if (!Schema_) {
+        builder->AppendString(ToString(nullptr));
+    } else if (Limit_ > 0 && Schema_->GetMemoryUsage(Limit_) <= Limit_) {
+        builder->AppendFormat("%v", *Schema_);
+    } else {
+        builder->AppendString("<schema memory usage is over the logging threshold>");
+    }
+}
+
+TFormatterWrapper<TTableSchemaTruncatedFormatter> MakeTableSchemaTruncatedFormatter(
+    const TTableSchemaPtr& schema,
+    i64 memoryLimit)
+{
+    return {TTableSchemaTruncatedFormatter(schema, memoryLimit)};
+}
+
+////////////////////////////////////////////////////////////////////////////////
+
 bool operator==(const TColumnSchema& lhs, const TColumnSchema& rhs)
 {
     return
diff --git a/yt/yt/client/table_client/schema.h b/yt/yt/client/table_client/schema.h
index 7cbbc86048..a74ed5f5d7 100644
--- a/yt/yt/client/table_client/schema.h
+++ b/yt/yt/client/table_client/schema.h
@@ -166,6 +166,7 @@ public:
     EValueType GetWireType() const;
 
     i64 GetMemoryUsage() const;
+    i64 GetMemoryUsage(i64 limit) const;
 
     // Check if column has plain old v1 type.
     bool IsOfV1Type() const;
@@ -398,6 +399,8 @@ public:
 
     i64 GetMemoryUsage() const;
 
+    i64 GetMemoryUsage(i64 limit) const;
+
 private:
     struct TColumnInfo
     {
@@ -462,6 +465,25 @@ void PrintTo(const TTableSchema& tableSchema, std::ostream* os);
 
 ////////////////////////////////////////////////////////////////////////////////
 
+class TTableSchemaTruncatedFormatter
+{
+public:
+    // NB: #schema is allowed to be |nullptr|.
+    TTableSchemaTruncatedFormatter(const TTableSchemaPtr& schema, i64 memoryLimit);
+
+    void operator()(TStringBuilderBase* builder) const;
+
+private:
+    const TTableSchema* const Schema_ = nullptr;
+    const i64 Limit_ = 0;
+};
+
+TFormatterWrapper<TTableSchemaTruncatedFormatter> MakeTableSchemaTruncatedFormatter(
+    const TTableSchemaPtr& schema,
+    i64 memoryLimit);
+
+////////////////////////////////////////////////////////////////////////////////
+
 bool operator == (const TColumnSchema& lhs, const TColumnSchema& rhs);
 
 bool operator == (const TDeletedColumn& lhs, const TDeletedColumn& rhs);