Merge pull request ClickHouse#114146 from Avogar/respect-subcolumn-sizes-setting-in-prewhere

Respect allow_calculating_subcolumns_sizes_for_merge_tree_reading in PREWHERE optimization

Author: Avogar

src/Interpreters/InterpreterSelectQuery.cpp

@@ -135,6 +135,7 @@ namespace Setting
     extern const SettingsMap additional_table_filters;
     extern const SettingsUInt64 aggregation_in_order_max_block_bytes;
     extern const SettingsUInt64 aggregation_memory_efficient_merge_threads;
+    extern const SettingsBool allow_calculating_subcolumns_sizes_for_merge_tree_reading;
     extern const SettingsUInt64 allow_experimental_parallel_reading_from_replicas;
     extern const SettingsUInt64 automatic_parallel_replicas_mode;
     extern const SettingsBool async_socket_for_remote;
@@ -876,7 +877,9 @@ InterpreterSelectQuery::InterpreterSelectQuery(
         {
             /// PREWHERE optimization: transfer some condition from WHERE to PREWHERE if enabled and viable
             Names queried_columns = syntax_analyzer_result->requiredSourceColumns();
-            if (const auto & column_sizes = storage->getColumnSizes(queried_columns); !column_sizes.empty())
+            const auto & column_sizes = storage->getColumnSizes(
+                queried_columns, context->getSettingsRef()[Setting::allow_calculating_subcolumns_sizes_for_merge_tree_reading]);
+            if (!column_sizes.empty())
             {
                 /// Extract column compressed sizes.
                 std::unordered_map<std::string, UInt64> column_compressed_sizes;

src/Processors/QueryPlan/Optimizations/optimizePrewhere.cpp

@@ -22,6 +22,7 @@ namespace DB
 {
 namespace Setting
 {
+    extern const SettingsBool allow_calculating_subcolumns_sizes_for_merge_tree_reading;
     extern const SettingsBool optimize_move_to_prewhere;
     extern const SettingsBool optimize_move_to_prewhere_if_final;
     extern const SettingsBool optimize_prewhere_after_pushdown;
@@ -224,7 +225,8 @@ void optimizePrewhere(QueryPlan::Node & parent_node, const bool remove_unused_co
 
     const auto & queried_columns = source_step_with_filter->requiredSourceColumns();
 
-    auto column_sizes = storage.getColumnSizes(queried_columns);
+    auto column_sizes = storage.getColumnSizes(
+        queried_columns, settings[Setting::allow_calculating_subcolumns_sizes_for_merge_tree_reading]);
     if (column_sizes.empty())
         return;
 

src/Storages/IStorage.h

@@ -216,9 +216,12 @@ class IStorage : public std::enable_shared_from_this<IStorage>, public TypePromo
     using ColumnSizeByName = std::unordered_map<std::string, ColumnSize>;
     virtual ColumnSizeByName getColumnSizes() const { return {}; }
 
-    /// Same as parameterless overload but also includes sizes for requested subcolumns
+    /// Same as parameterless overload but also includes sizes for the requested subcolumns.
+    /// Computing exact subcolumn sizes can be expensive, so `calculate_subcolumn_sizes` (driven by
+    /// `allow_calculating_subcolumns_sizes_for_merge_tree_reading` at call sites) selects between the
+    /// exact size and the cheaper top-level column size as an approximation.
     /// The default implementation falls back to the parameterless version.
-    virtual ColumnSizeByName getColumnSizes(const Names & /*columns*/) const { return getColumnSizes(); }
+    virtual ColumnSizeByName getColumnSizes(const Names & /*columns*/, bool /*calculate_subcolumn_sizes*/) const { return getColumnSizes(); }
 
     /// Same as getColumnSizes() but may return nullopt in some specific engines like Merge/Alias
     virtual std::optional<ColumnSizeByName> tryGetColumnSizes() const { return getColumnSizes(); }

src/Storages/MergeTree/MergeTreeData.cpp

@@ -7760,7 +7760,7 @@ void MergeTreeData::addPartContributionToColumnAndSecondaryIndexSizesUnlocked(co
     primary_index_size.add(part->getIndexSizeFromFile());
 }
 
-IStorage::ColumnSizeByName MergeTreeData::getColumnSizes(const Names & columns) const
+IStorage::ColumnSizeByName MergeTreeData::getColumnSizes(const Names & columns, bool calculate_subcolumn_sizes) const
 {
     auto result = getColumnSizes();
 
@@ -7777,12 +7777,36 @@ IStorage::ColumnSizeByName MergeTreeData::getColumnSizes(const Names & columns)
     if (subcolumn_names.empty())
         return result;
 
-    /// For each requested column that is a subcolumn and not already in the result,
-    /// aggregate its size across all active parts using getSubcolumnSize.
-    /// This gives the correct on-disk size for subcolumns based on required substreams.
-    auto parts_lock = readLockParts();
-    auto committed_parts_range = getDataPartsStateRange(DataPartState::Active);
-    for (const auto & part : committed_parts_range)
+    /// When exact subcolumn sizes are disabled, approximate each subcolumn with its whole top-level
+    /// column size (already in result), resolved from metadata without locking parts.
+    if (!calculate_subcolumn_sizes)
+    {
+        auto metadata_snapshot = getInMemoryMetadataPtr(getContext(), false);
+        const auto & storage_columns = metadata_snapshot->getColumns();
+        for (const auto & col_name : subcolumn_names)
+        {
+            auto column = storage_columns.tryGetColumnOrSubcolumn(GetColumnsOptions::All, col_name);
+            if (column && column->isSubcolumn())
+            {
+                if (auto it = result.find(column->getNameInStorage()); it != result.end())
+                    result[col_name] = it->second;
+            }
+        }
+
+        return result;
+    }
+
+    /// Exact subcolumn sizes are derived per active part from the required substreams.
+    /// Snapshot the parts under the lock and release it before the per-part size calculation,
+    /// which reads part-local state and would otherwise block part commits and merges.
+    DataPartsVector parts;
+    {
+        auto parts_lock = readLockParts();
+        auto committed_parts_range = getDataPartsStateRange(DataPartState::Active);
+        parts.assign(committed_parts_range.begin(), committed_parts_range.end());
+    }
+
+    for (const auto & part : parts)
     {
         for (const auto & col_name : subcolumn_names)
         {

src/Storages/MergeTree/MergeTreeData.h

@@ -1158,7 +1158,7 @@ class MergeTreeData : public WithMutableContext, public IStorage, public IBackgr
         return column_sizes;
     }
 
-    ColumnSizeByName getColumnSizes(const Names & columns) const override;
+    ColumnSizeByName getColumnSizes(const Names & columns, bool calculate_subcolumn_sizes) const override;
 
     IndexSizeByName getSecondaryIndexSizes() const override
     {

src/Storages/PostgreSQL/StorageMaterializedPostgreSQL.cpp

@@ -354,10 +354,10 @@ IStorage::ColumnSizeByName StorageMaterializedPostgreSQL::getColumnSizes() const
 }
 
 
-IStorage::ColumnSizeByName StorageMaterializedPostgreSQL::getColumnSizes(const Names & columns) const
+IStorage::ColumnSizeByName StorageMaterializedPostgreSQL::getColumnSizes(const Names & columns, bool calculate_subcolumn_sizes) const
 {
     if (auto nested = tryGetNested())
-        return nested->getColumnSizes(columns);
+        return nested->getColumnSizes(columns, calculate_subcolumn_sizes);
     return {};
 }
 

src/Storages/PostgreSQL/StorageMaterializedPostgreSQL.h

@@ -107,7 +107,7 @@ class StorageMaterializedPostgreSQL final : public IStorage, WithContext
     bool supportsOptimizationToSubcolumns() const override;
 
     ColumnSizeByName getColumnSizes() const override;
-    ColumnSizeByName getColumnSizes(const Names & columns) const override;
+    ColumnSizeByName getColumnSizes(const Names & columns, bool calculate_subcolumn_sizes) const override;
 
     std::optional<UInt64> totalRows(ContextPtr query_context) const override;
     std::optional<UInt64> totalBytes(ContextPtr query_context) const override;

src/Storages/StorageAlias.h

@@ -270,7 +270,7 @@ class StorageAlias final : public IStorage, WithContext
     }
 
     ColumnSizeByName getColumnSizes() const override { auto target = tryGetTargetTable(); return target ? target->getColumnSizes() : ColumnSizeByName{}; }
-    ColumnSizeByName getColumnSizes(const Names & columns) const override { auto target = tryGetTargetTable(); return target ? target->getColumnSizes(columns) : ColumnSizeByName{}; }
+    ColumnSizeByName getColumnSizes(const Names & columns, bool calculate_subcolumn_sizes) const override { auto target = tryGetTargetTable(); return target ? target->getColumnSizes(columns, calculate_subcolumn_sizes) : ColumnSizeByName{}; }
     std::optional<ColumnSizeByName> tryGetColumnSizes() const override
     {
         auto target = tryGetTargetTable();

src/Storages/StorageMerge.cpp

@@ -1986,13 +1986,13 @@ IStorage::ColumnSizeByName StorageMerge::getColumnSizes() const
     return column_sizes;
 }
 
-IStorage::ColumnSizeByName StorageMerge::getColumnSizes(const Names & columns) const
+IStorage::ColumnSizeByName StorageMerge::getColumnSizes(const Names & columns, bool calculate_subcolumn_sizes) const
 {
     ColumnSizeByName column_sizes;
 
     forEachTable([&](const auto & table)
     {
-        for (const auto & [name, size] : table->getColumnSizes(columns))
+        for (const auto & [name, size] : table->getColumnSizes(columns, calculate_subcolumn_sizes))
             column_sizes[name].add(size);
     });
 

src/Storages/StorageMerge.h

@@ -151,7 +151,7 @@ class StorageMerge final : public IStorage, WithContext
         const IStorage * ignore_self);
 
     ColumnSizeByName getColumnSizes() const override;
-    ColumnSizeByName getColumnSizes(const Names & columns) const override;
+    ColumnSizeByName getColumnSizes(const Names & columns, bool calculate_subcolumn_sizes) const override;
 
     std::optional<ColumnSizeByName> tryGetColumnSizes() const override;