[fix](be) Count variant data in information_schema.column_data_sizes (#68444)

be/src/information_schema/schema_column_data_sizes_scanner.cpp

@@ -111,40 +111,57 @@ void SchemaColumnDataSizesScanner::_collect_column_data_sizes_from_rowsets(
 
         // Get column data page stats from each segment footer and aggregate by column_unique_id
         for (const auto& segment : segments) {
-            auto collector = [&](const segment_v2::ColumnMetaPB& column_meta) {
+            auto collector = [&](const segment_v2::ColumnMetaPB& column_meta,
+                                 uint32_t column_unique_id) {
                 if (column_meta.has_compressed_data_bytes() &&
                     column_meta.has_uncompressed_data_bytes()) {
-                    auto cid = schema->field_index(column_meta.unique_id());
+                    auto cid = schema->field_index(column_unique_id);
                     if (cid == -1) {
                         return;
                     }
-                    // Aggregate stats by column_unique_id
-                    if (aggregated_stats.contains(column_meta.unique_id())) {
-                        auto& existing_stats = aggregated_stats[column_meta.unique_id()];
-                        existing_stats.compressed_data_bytes += column_meta.compressed_data_bytes();
-                        existing_stats.uncompressed_data_bytes +=
-                                column_meta.uncompressed_data_bytes();
-                        existing_stats.raw_data_bytes += column_meta.raw_data_bytes();
-                    } else {
-                        aggregated_stats[column_meta.unique_id()] = ColumnDataSizeInfo {
+                    auto [it, inserted] = aggregated_stats.try_emplace(column_unique_id);
+                    auto& stats = it->second;
+                    if (inserted) {
+                        stats = ColumnDataSizeInfo {
                                 .backend_id = backend_id_,
                                 .table_id = table_id,
                                 .index_id = index_id,
                                 .partition_id = partition_id,
                                 .tablet_id = tablet_id,
                                 .rowset_id = rowset_id,
-                                .column_unique_id = column_meta.unique_id(),
+                                .column_unique_id = column_unique_id,
                                 .column_name = schema->column(cid).name(),
                                 .column_type = TabletColumn::get_string_by_field_type(
-                                        static_cast<FieldType>(column_meta.type())),
-                                .compressed_data_bytes = column_meta.compressed_data_bytes(),
-                                .uncompressed_data_bytes = column_meta.uncompressed_data_bytes(),
-                                .raw_data_bytes = column_meta.raw_data_bytes(),
+                                        schema->column(cid).type()),
+                                .compressed_data_bytes = 0,
+                                .uncompressed_data_bytes = 0,
+                                .raw_data_bytes = 0,
                         };
                     }
+                    stats.compressed_data_bytes += column_meta.compressed_data_bytes();
+                    stats.uncompressed_data_bytes += column_meta.uncompressed_data_bytes();
+                    stats.raw_data_bytes += column_meta.raw_data_bytes();
+                }
+            };
+            auto footer_collector = [&](const segment_v2::ColumnMetaPB& column_meta) {
+                // A variant meta counts the data of all its inner writers, whose own metas carry no
+                // data bytes. Compaction writes variant subcolumns as columns of their own, with
+                // data bytes but no unique id; count them to their root variant column.
+                uint32_t column_unique_id = column_meta.unique_id();
+                if (column_meta.unique_id() == -1 && column_meta.has_compressed_data_bytes()) {
+                    DORIS_CHECK(column_meta.has_column_path_info());
+                    column_unique_id = column_meta.column_path_info().parrent_column_unique_id();
+                }
+                collector(column_meta, column_unique_id);
+                if (column_meta.type() == static_cast<int>(FieldType::OLAP_FIELD_TYPE_VARIANT)) {
+                    // External column meta (V3) embeds the sparse and doc value columns of a
+                    // variant into its root column meta.
+                    for (const auto& child : column_meta.children_columns()) {
+                        collector(child, column_unique_id);
+                    }
                 }
             };
-            st = segment->traverse_column_meta_pbs(collector);
+            st = segment->traverse_column_meta_pbs(footer_collector);
             if (!st.ok()) {
                 continue;
             }
@@ -354,4 +371,4 @@ Status SchemaColumnDataSizesScanner::_fill_block_impl(Block* block) {
     return Status::OK();
 }
 
-} // namespace doris
+} // namespace doris

be/src/storage/merger.cpp

@@ -706,8 +706,9 @@ Status Merger::vertical_merge_rowsets(
             const auto& col = tablet_schema.column(col_ordinal);
             int32_t uid = col.unique_id();
 
-            // Variant columns (root or subcolumn): raw_data_bytes is 0 (TODO in writer),
-            // cannot estimate from footer, fallback to default for the entire group.
+            // Variant columns (root or subcolumn): subcolumns written as columns of their own
+            // (e.g. by compaction) are not counted in the root's raw_data_bytes, cannot estimate
+            // from footer, fallback to default for the entire group.
             if (uid < 0 || col.is_variant_type()) {
                 need_fallback = true;
                 break;

be/src/storage/segment/column_writer.cpp

@@ -1342,6 +1342,19 @@ uint64_t VariantColumnWriter::estimate_buffer_size() {
     return _impl->estimate_buffer_size();
 }
 
+uint64_t VariantColumnWriter::get_raw_data_bytes() const {
+    return _impl->get_total_data_pages_bytes(&ColumnWriter::get_raw_data_bytes);
+}
+
+uint64_t VariantColumnWriter::get_total_uncompressed_data_pages_bytes() const {
+    return _impl->get_total_data_pages_bytes(
+            &ColumnWriter::get_total_uncompressed_data_pages_bytes);
+}
+
+uint64_t VariantColumnWriter::get_total_compressed_data_pages_bytes() const {
+    return _impl->get_total_data_pages_bytes(&ColumnWriter::get_total_compressed_data_pages_bytes);
+}
+
 Status VariantColumnWriter::finish() {
     return _impl->finish();
 }

be/src/storage/segment/column_writer.h

@@ -210,6 +210,8 @@ public:
     virtual uint64_t get_raw_data_bytes() const = 0;
     virtual uint64_t get_total_uncompressed_data_pages_bytes() const = 0;
     virtual uint64_t get_total_compressed_data_pages_bytes() const = 0;
+    // Points to one of the three data bytes getters above.
+    using DataBytesGetter = uint64_t (ColumnWriter::*)() const;
 
     // used for append not null data.
     virtual Status append_data(const uint8_t** ptr, size_t num_rows) = 0;
@@ -620,16 +622,14 @@ public:
     Status write_bloom_filter_index() override;
     ordinal_t get_next_rowid() const override { return _next_rowid; }
 
-    uint64_t get_raw_data_bytes() const override {
-        return 0; // TODO
-    }
+    uint64_t get_raw_data_bytes() const override { return _writer->get_raw_data_bytes(); }
 
     uint64_t get_total_uncompressed_data_pages_bytes() const override {
-        return 0; // TODO
+        return _writer->get_total_uncompressed_data_pages_bytes();
     }
 
     uint64_t get_total_compressed_data_pages_bytes() const override {
-        return 0; // TODO
+        return _writer->get_total_compressed_data_pages_bytes();
     }
 
     Status append_nulls(size_t num_rows) override {
@@ -688,17 +688,11 @@ public:
     Status write_bloom_filter_index() override;
     ordinal_t get_next_rowid() const override { return _next_rowid; }
 
-    uint64_t get_raw_data_bytes() const override {
-        return 0; // TODO
-    }
+    uint64_t get_raw_data_bytes() const override;
 
-    uint64_t get_total_uncompressed_data_pages_bytes() const override {
-        return 0; // TODO
-    }
+    uint64_t get_total_uncompressed_data_pages_bytes() const override;
 
-    uint64_t get_total_compressed_data_pages_bytes() const override {
-        return 0; // TODO
-    }
+    uint64_t get_total_compressed_data_pages_bytes() const override;
 
     Status append_nulls(size_t num_rows) override {
         return Status::NotSupported("variant writer can not append_nulls");

be/src/storage/segment/variant/v2/variant_column_writer.cpp

@@ -359,4 +359,17 @@ uint64_t VariantV2ColumnWriter::estimate_buffer_size() {
     return size;
 }
 
+uint64_t VariantV2ColumnWriter::get_total_data_pages_bytes(
+        ColumnWriter::DataBytesGetter getter) const {
+    DORIS_CHECK(_is_finalized);
+    uint64_t size = std::invoke(getter, _root_writer.get());
+    for (const auto& writer : _subcolumn_writers) {
+        size += std::invoke(getter, writer.get());
+    }
+    if (_binary_writer != nullptr) {
+        size += _binary_writer->get_total_data_pages_bytes(getter);
+    }
+    return size;
+}
+
 } // namespace doris::segment_v2

be/src/storage/segment/variant/v2/variant_column_writer.h

@@ -59,6 +59,9 @@ public:
     Status write_inverted_index();
     Status write_bloom_filter_index();
     uint64_t estimate_buffer_size();
+    // Sums `getter`, one of the ColumnWriter data bytes getters, over every column writer of this
+    // variant: the root, the materialized subcolumns and the sparse or doc value columns.
+    uint64_t get_total_data_pages_bytes(ColumnWriter::DataBytesGetter getter) const;
 
 private:
     Status _write_root(const IColumn* root_jsonb, int& column_id);

be/src/storage/segment/variant/variant_column_writer_impl.cpp

@@ -211,6 +211,18 @@ uint64_t UnifiedSparseColumnWriter::estimate_buffer_size() const {
     return size;
 }
 
+uint64_t UnifiedSparseColumnWriter::get_total_data_pages_bytes(
+        ColumnWriter::DataBytesGetter getter) const {
+    uint64_t size = 0;
+    if (_single_writer) {
+        size += std::invoke(getter, _single_writer.get());
+    }
+    for (const auto& w : _bucket_writers) {
+        size += std::invoke(getter, w.get());
+    }
+    return size;
+}
+
 Status UnifiedSparseColumnWriter::finish() {
     if (_single_writer) {
         RETURN_IF_ERROR(_single_writer->finish());
@@ -350,6 +362,17 @@ uint64_t VariantDocWriter::estimate_buffer_size() const {
     return size;
 }
 
+uint64_t VariantDocWriter::get_total_data_pages_bytes(ColumnWriter::DataBytesGetter getter) const {
+    uint64_t size = 0;
+    for (const auto& writer : _subcolumn_writers) {
+        size += std::invoke(getter, writer.get());
+    }
+    for (const auto& writer : _doc_value_column_writers) {
+        size += std::invoke(getter, writer.get());
+    }
+    return size;
+}
+
 Status VariantDocWriter::finish() {
     for (auto& writer : _subcolumn_writers) {
         RETURN_IF_ERROR(writer->finish());
@@ -520,6 +543,11 @@ uint64_t VariantColumnWriterImpl::estimate_buffer_size() {
     return _v2_writer ? _v2_writer->estimate_buffer_size() : 0;
 }
 
+uint64_t VariantColumnWriterImpl::get_total_data_pages_bytes(
+        ColumnWriter::DataBytesGetter getter) const {
+    return _v2_writer->get_total_data_pages_bytes(getter);
+}
+
 VariantSubcolumnWriter::VariantSubcolumnWriter(const ColumnWriterOptions& opts,
                                                TabletColumnPtr column)
         : ColumnWriter(std::move(column), opts.meta->is_nullable(), opts.meta) {

be/src/storage/segment/variant/variant_column_writer_impl.h

@@ -62,6 +62,8 @@ public:
     virtual Status write_inverted_index() = 0;
     virtual Status write_bloom_filter_index() = 0;
     virtual uint64_t estimate_buffer_size() const = 0;
+    // Sums `getter`, one of the ColumnWriter data bytes getters, over the owned column writers.
+    virtual uint64_t get_total_data_pages_bytes(ColumnWriter::DataBytesGetter getter) const = 0;
     virtual void merge_stats_to(VariantStatistics* stats) const = 0;
 };
 
@@ -80,6 +82,7 @@ public:
     Status write_inverted_index() override;
     Status write_bloom_filter_index() override;
     uint64_t estimate_buffer_size() const override;
+    uint64_t get_total_data_pages_bytes(ColumnWriter::DataBytesGetter getter) const override;
     void merge_stats_to(VariantStatistics* stats) const override;
 
 private:
@@ -127,6 +130,7 @@ public:
     Status write_zone_map() override;
     Status write_inverted_index() override;
     Status write_bloom_filter_index() override;
+    uint64_t get_total_data_pages_bytes(ColumnWriter::DataBytesGetter getter) const override;
     void merge_stats_to(VariantStatistics* stats) const override;
 
 private:
@@ -171,6 +175,7 @@ public:
     Status write_inverted_index();
     Status write_bloom_filter_index();
     uint64_t estimate_buffer_size();
+    uint64_t get_total_data_pages_bytes(ColumnWriter::DataBytesGetter getter) const;
 
 private:
     Status _ensure_writer();
@@ -205,15 +210,15 @@ public:
     ordinal_t get_next_rowid() const override { return _next_rowid; }
 
     uint64_t get_raw_data_bytes() const override {
-        return 0; // TODO
+        return _get_total_data_pages_bytes(&ColumnWriter::get_raw_data_bytes);
     }
 
     uint64_t get_total_uncompressed_data_pages_bytes() const override {
-        return 0; // TODO
+        return _get_total_data_pages_bytes(&ColumnWriter::get_total_uncompressed_data_pages_bytes);
     }
 
     uint64_t get_total_compressed_data_pages_bytes() const override {
-        return 0; // TODO
+        return _get_total_data_pages_bytes(&ColumnWriter::get_total_compressed_data_pages_bytes);
     }
 
     Status append_nulls(size_t num_rows) override {
@@ -228,6 +233,15 @@ public:
     Status finalize();
 
 private:
+    template <typename Func>
+    uint64_t _get_total_data_pages_bytes(Func func) const {
+        uint64_t size = std::invoke(func, _doc_value_column_writer.get());
+        for (const auto& writer : _subcolumn_writers) {
+            size += std::invoke(func, writer.get());
+        }
+        return size;
+    }
+
     Status _append(const uint8_t* null_map, const uint8_t** ptr, size_t num_rows);
     Status _ensure_input_format(const VariantColumnData& column);
     Status _initialize_v2_shredder();

be/test/storage/variant/variant_column_writer_reader_test.cpp

@@ -1704,6 +1704,8 @@ protected:
 
         RETURN_IF_ERROR(writer->finish());
         RETURN_IF_ERROR(writer->write_data());
+        // The file holds only the data pages of this variant so far.
+        EXPECT_EQ(writer->get_total_compressed_data_pages_bytes(), file_writer->bytes_appended());
         RETURN_IF_ERROR(writer->write_ordinal_index());
         RETURN_IF_ERROR(writer->write_zone_map());
         if (index_write_policy == VariantIndexWritePolicy::BLOOM_AND_INVERTED) {
@@ -1775,6 +1777,8 @@ protected:
 
         RETURN_IF_ERROR(writer->finish());
         RETURN_IF_ERROR(writer->write_data());
+        // The file holds only the data pages of this subcolumn so far.
+        EXPECT_EQ(writer->get_total_compressed_data_pages_bytes(), file_writer->bytes_appended());
         RETURN_IF_ERROR(writer->write_ordinal_index());
         RETURN_IF_ERROR(writer->write_zone_map());
         RETURN_IF_ERROR(file_writer->close());
@@ -3951,6 +3955,10 @@ TEST_F(VariantColumnWriterReaderTest, test_write_data_normal) {
     EXPECT_TRUE(st.ok()) << st.msg();
     st = writer->write_data();
     EXPECT_TRUE(st.ok()) << st.msg();
+    // The root, the subcolumns and the sparse columns all count to the variant.
+    EXPECT_EQ(writer->get_total_compressed_data_pages_bytes(), file_writer->bytes_appended());
+    EXPECT_GT(writer->get_total_uncompressed_data_pages_bytes(), 0);
+    EXPECT_GT(writer->get_raw_data_bytes(), 0);
     st = writer->write_ordinal_index();
     EXPECT_TRUE(st.ok()) << st.msg();
     st = writer->write_zone_map();
@@ -5241,6 +5249,12 @@ TEST_P(VariantSpecializedWriterCompatibilityTest, doc_compact_writer_round_trip)
     EXPECT_TRUE(doc_compact_writer->finish().ok());
     EXPECT_TRUE(root_writer->write_data().ok());
     EXPECT_TRUE(doc_compact_writer->write_data().ok());
+    // The file holds only data pages so far: the root of the root-only writer, plus the doc value
+    // column and the materialized subcolumns of the doc compact writer.
+    EXPECT_GT(doc_compact_writer->get_total_compressed_data_pages_bytes(), 0);
+    EXPECT_EQ(root_writer->get_total_compressed_data_pages_bytes() +
+                      doc_compact_writer->get_total_compressed_data_pages_bytes(),
+              file_writer->bytes_appended());
     EXPECT_TRUE(root_writer->write_ordinal_index().ok());
     EXPECT_TRUE(doc_compact_writer->write_ordinal_index().ok());
     EXPECT_TRUE(file_writer->close().ok());

regression-test/data/query_p0/system/test_query_sys_column_data_sizes_variant.out

@@ -0,0 +1,33 @@
+-- This file is automatically generated. You should know what you did if you want to edit this
+-- !sizes --
+id	BIGINT	true	true	true
+v	VARIANT	true	true	true
+
+-- !sizes --
+id	BIGINT	true	true	true
+v	VARIANT	true	true	true
+
+-- !sizes --
+id	BIGINT	true	true	true
+v	VARIANT	true	true	true
+
+-- !sizes --
+id	BIGINT	true	true	true
+v	VARIANT	true	true	true
+
+-- !sizes --
+id	BIGINT	true	true	true
+v	VARIANT	true	true	true
+
+-- !sizes --
+id	BIGINT	true	true	true
+v	VARIANT	true	true	true
+
+-- !sizes --
+id	BIGINT	true	true	true
+v	VARIANT	true	true	true
+
+-- !sizes --
+id	BIGINT	true	true	true
+v	VARIANT	true	true	true
+

regression-test/suites/query_p0/system/test_query_sys_column_data_sizes_variant.groovy

@@ -0,0 +1,120 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements.  See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership.  The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License.  You may obtain a copy of the License at
+//
+//   http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied.  See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+suite("test_query_sys_column_data_sizes_variant", "query,p0") {
+    // A variant column keeps its data in the root, the materialized subcolumns and the sparse or
+    // doc value columns. Compaction writes them as columns of their own, and external column meta
+    // (V3) embeds the sparse and doc value columns into the root meta. All of them count to the
+    // variant column.
+    def checkVariantDataSizes = { String tableName ->
+        for (int i = 0; i < 3; i++) {
+            sql """
+                INSERT INTO ${tableName}
+                SELECT number, parse_to_variant(concat('{"a":', number, ',"b":"', repeat('x', number % 32),
+                        '","c":{"d":', number * 2, '},"k', number % 16, '":"', repeat('y', number % 64), '"}'))
+                FROM numbers("number" = "10000")
+            """
+        }
+        def tabletId = sql("SHOW TABLETS FROM ${tableName}")[0][0]
+        def check = {
+            qt_sizes """
+                SELECT COLUMN_NAME, COLUMN_TYPE, SUM(COMPRESSED_DATA_BYTES) > 0,
+                       SUM(UNCOMPRESSED_DATA_BYTES) > 0, SUM(RAW_DATA_BYTES) > 0
+                FROM information_schema.column_data_sizes
+                WHERE TABLET_ID = ${tabletId}
+                GROUP BY COLUMN_NAME, COLUMN_TYPE
+                ORDER BY COLUMN_NAME
+            """
+            // The data pages of all columns lie in the segment files and make up most of them, so
+            // a part of the variant that is not counted, or counted twice, breaks the bounds.
+            def bytes = sql """
+                SELECT SUM(c.data_page_bytes), SUM(r.DATA_DISK_SIZE)
+                FROM (
+                    SELECT BACKEND_ID, ROWSET_ID, SUM(COMPRESSED_DATA_BYTES) AS data_page_bytes
+                    FROM information_schema.column_data_sizes
+                    WHERE TABLET_ID = ${tabletId}
+                    GROUP BY BACKEND_ID, ROWSET_ID
+                ) c
+                JOIN information_schema.rowsets r
+                    ON r.BACKEND_ID = c.BACKEND_ID AND r.ROWSET_ID = c.ROWSET_ID
+            """
+            long dataPageBytes = bytes[0][0] as long
+            long segmentBytes = bytes[0][1] as long
+            logger.info("${tableName}: ${dataPageBytes} data page bytes in ${segmentBytes} segment bytes")
+            assertTrue(dataPageBytes <= segmentBytes && dataPageBytes * 2 > segmentBytes,
+                    "${tableName}: ${dataPageBytes} data page bytes in ${segmentBytes} segment bytes")
+        }
+
+        check()
+        trigger_and_wait_compaction(tableName, "full")
+        check()
+    }
+
+    // Materialized subcolumns and bucketized sparse columns, with the column metas in the footer
+    // (V2) and in the external column meta region (V3).
+    sql "set default_variant_enable_doc_mode = false"
+    sql "DROP TABLE IF EXISTS test_column_data_sizes_variant_v2"
+    sql """
+        CREATE TABLE test_column_data_sizes_variant_v2 (
+            id BIGINT,
+            v VARIANT<PROPERTIES("variant_max_subcolumns_count" = "1", "variant_sparse_hash_shard_count" = "2")>
+        )
+        DUPLICATE KEY(id)
+        DISTRIBUTED BY HASH(id) BUCKETS 1
+        PROPERTIES ("replication_num" = "1", "disable_auto_compaction" = "true", "storage_format" = "V2")
+    """
+    checkVariantDataSizes("test_column_data_sizes_variant_v2")
+
+    sql "DROP TABLE IF EXISTS test_column_data_sizes_variant_v3"
+    sql """
+        CREATE TABLE test_column_data_sizes_variant_v3 (
+            id BIGINT,
+            v VARIANT<PROPERTIES("variant_max_subcolumns_count" = "1", "variant_sparse_hash_shard_count" = "2")>
+        )
+        DUPLICATE KEY(id)
+        DISTRIBUTED BY HASH(id) BUCKETS 1
+        PROPERTIES ("replication_num" = "1", "disable_auto_compaction" = "true", "storage_format" = "V3")
+    """
+    checkVariantDataSizes("test_column_data_sizes_variant_v3")
+
+    // Doc value columns and materialized subcolumns in doc mode.
+    sql "DROP TABLE IF EXISTS test_column_data_sizes_variant_doc_v2"
+    sql """
+        CREATE TABLE test_column_data_sizes_variant_doc_v2 (
+            id BIGINT,
+            v VARIANT<PROPERTIES("variant_enable_doc_mode" = "true", "variant_doc_materialization_min_rows" = "0",
+                    "variant_doc_hash_shard_count" = "2")>
+        )
+        DUPLICATE KEY(id)
+        DISTRIBUTED BY HASH(id) BUCKETS 1
+        PROPERTIES ("replication_num" = "1", "disable_auto_compaction" = "true", "storage_format" = "V2")
+    """
+    checkVariantDataSizes("test_column_data_sizes_variant_doc_v2")
+
+    sql "DROP TABLE IF EXISTS test_column_data_sizes_variant_doc_v3"
+    sql """
+        CREATE TABLE test_column_data_sizes_variant_doc_v3 (
+            id BIGINT,
+            v VARIANT<PROPERTIES("variant_enable_doc_mode" = "true", "variant_doc_materialization_min_rows" = "0",
+                    "variant_doc_hash_shard_count" = "2")>
+        )
+        DUPLICATE KEY(id)
+        DISTRIBUTED BY HASH(id) BUCKETS 1
+        PROPERTIES ("replication_num" = "1", "disable_auto_compaction" = "true", "storage_format" = "V3")
+    """
+    checkVariantDataSizes("test_column_data_sizes_variant_doc_v3")
+}