[Feature] table sample

be/src/exec/pipeline/scan/olap_chunk_source.cpp

@@ -236,6 +236,7 @@ Status OlapChunkSource::_init_reader_params(const std::vector<std::unique_ptr<Ol
     _params.runtime_state = _runtime_state;
     _params.use_page_cache = _runtime_state->use_page_cache();
     _params.use_pk_index = thrift_olap_scan_node.use_pk_index;
+    _params.sample_options = thrift_olap_scan_node.sample_options;
     if (thrift_olap_scan_node.__isset.enable_prune_column_after_index_filter) {
         _params.prune_column_after_index_filter = thrift_olap_scan_node.enable_prune_column_after_index_filter;
     }
@@ -778,6 +779,23 @@ void OlapChunkSource::_update_counter() {
         RuntimeProfile::Counter* c = ADD_CHILD_TIMER(_runtime_profile, "FlatJsonFlatten", parent_name);
         COUNTER_UPDATE(c, _reader->stats().json_flatten_ns);
     }
+
+    // Data sampling
+    if (_params.sample_options.enable_sampling) {
+        _runtime_profile->add_info_string("SampleMethod", to_string(_params.sample_options.sample_method));
+        _runtime_profile->add_info_string("SamplePercent",
+                                          std::to_string(_params.sample_options.probability_percent) + "%");
+        COUNTER_UPDATE(ADD_CHILD_TIMER(_runtime_profile, "SampleTime", parent_name),
+                       _reader->stats().sample_population_size);
+        COUNTER_UPDATE(ADD_CHILD_TIMER(_runtime_profile, "SampleBuildHistogramTime", parent_name),
+                       _reader->stats().sample_build_histogram_time_ns);
+        COUNTER_UPDATE(ADD_CHILD_COUNTER(_runtime_profile, "SampleSize", TUnit::UNIT, parent_name),
+                       _reader->stats().sample_size);
+        COUNTER_UPDATE(ADD_CHILD_COUNTER(_runtime_profile, "SamplePopulationSize", TUnit::UNIT, parent_name),
+                       _reader->stats().sample_population_size);
+        COUNTER_UPDATE(ADD_CHILD_COUNTER(_runtime_profile, "SampleBuildHistogramCount", TUnit::UNIT, parent_name),
+                       _reader->stats().sample_build_histogram_count);
+    }
 }
 
 } // namespace starrocks::pipeline

be/src/storage/CMakeLists.txt

@@ -99,8 +99,9 @@ set(STORAGE_FILES
     rowset/ordinal_page_index.cpp
     rowset/page_io.cpp
     rowset/binary_dict_page.cpp
-    rowset/dict_page.cpp
     rowset/binary_prefix_page.cpp
+    rowset/data_sample.cpp
+    rowset/dict_page.cpp
     rowset/segment.cpp
     rowset/segment_writer.cpp
     rowset/segment_rewriter.cpp

be/src/storage/olap_common.h

@@ -313,6 +313,13 @@ struct OlapReaderStatistics {
     std::unordered_map<std::string, int64_t> flat_json_hits;
     std::unordered_map<std::string, int64_t> merge_json_hits;
     std::unordered_map<std::string, int64_t> dynamic_json_hits;
+
+    // Counters for data sampling
+    int64_t sample_time_ns = 0;               // Records the time to prepare sample, actual IO time is not included
+    int64_t sample_size = 0;                  // Records the number of hits in the sample. Granularity can be BLOCK/PAGE
+    int64_t sample_population_size = 0;       // Records the total number of samples. Granularity can be BLOCK/PAGE
+    int64_t sample_build_histogram_count = 0; // Records the number of histogram built for sampling
+    int64_t sample_build_histogram_time_ns = 0; // Records the time to build histogram
 };
 
 // OlapWriterStatistics used to collect statistics when write data to storage

be/src/storage/range.h

@@ -489,11 +489,17 @@ inline std::ostream& operator<<(std::ostream& os, const SparseRange<T>& range) {
 
 template class Range<>;
 template class Range<ordinal_t>;
+using RowIdRange = Range<rowid_t>;
+using OridinalRange = Range<ordinal_t>;
 
 template class SparseRange<>;
 template class SparseRange<ordinal_t>;
+using RowIdSparseRange = SparseRange<rowid_t>;
+using OridinalSparseRange = SparseRange<ordinal_t>;
 
 template class SparseRangeIterator<>;
 template class SparseRangeIterator<ordinal_t>;
+using RowIdSparseRangeIterator = SparseRangeIterator<rowid_t>;
+using OrdinalSparseRangeIterator = SparseRangeIterator<ordinal_t>;
 
 } // namespace starrocks

be/src/storage/rowset/column_iterator.h

@@ -158,6 +158,8 @@ class ColumnIterator {
 
     virtual ordinal_t get_current_ordinal() const = 0;
 
+    virtual bool has_zone_map() const { return false; }
+
     /// Store the row ranges that satisfy the given predicates into |row_ranges|.
     /// |pred_relation| is the relation among |predicates|, it can be AND or OR.
     virtual Status get_row_ranges_by_zone_map(const std::vector<const ColumnPredicate*>& predicates,
@@ -259,9 +261,11 @@ class ColumnIterator {
 
     virtual Status null_count(size_t* count) { return Status::OK(); };
 
+    // RAW interface, should be used carefully
+    virtual ColumnReader* get_column_reader() { return nullptr; };
+
 protected:
     ColumnIteratorOptions _opts;
-    virtual ColumnReader* get_column_reader() { return nullptr; };
 };
 
 } // namespace starrocks

be/src/storage/rowset/column_reader.cpp

@@ -63,6 +63,7 @@
 #include "storage/rowset/encoding_info.h"
 #include "storage/rowset/json_column_iterator.h"
 #include "storage/rowset/map_column_iterator.h"
+#include "storage/rowset/options.h"
 #include "storage/rowset/page_handle.h"
 #include "storage/rowset/page_io.h"
 #include "storage/rowset/page_pointer.h"
@@ -559,6 +560,11 @@ Status ColumnReader::seek_by_page_index(int page_index, OrdinalPageIndexIterator
     return Status::OK();
 }
 
+std::pair<ordinal_t, ordinal_t> ColumnReader::get_page_range(size_t page_index) {
+    DCHECK(_ordinal_index);
+    return std::make_pair(_ordinal_index->get_first_ordinal(page_index), _ordinal_index->get_last_ordinal(page_index));
+}
+
 Status ColumnReader::zone_map_filter(const std::vector<const ColumnPredicate*>& predicates,
                                      const ColumnPredicate* del_predicate,
                                      std::unordered_set<uint32_t>* del_partial_filtered_pages,
@@ -579,6 +585,24 @@ Status ColumnReader::zone_map_filter(const std::vector<const ColumnPredicate*>&
     return Status::OK();
 }
 
+StatusOr<std::vector<ZoneMapDetail>> ColumnReader::get_raw_zone_map(const IndexReadOptions& opts) {
+    RETURN_IF_ERROR(_load_zonemap_index(opts));
+    DCHECK(_zonemap_index);
+    DCHECK(_zonemap_index->loaded());
+
+    LogicalType type = _encoding_info->type();
+    int32_t num_pages = _zonemap_index->num_pages();
+    std::vector<ZoneMapDetail> result(num_pages);
+
+    for (auto& zm : _zonemap_index->page_zone_maps()) {
+        ZoneMapDetail detail;
+        RETURN_IF_ERROR(_parse_zone_map(type, zm, &detail));
+        result.emplace_back(detail);
+    }
+
+    return result;
+}
+
 template <CompoundNodeType PredRelation>
 Status ColumnReader::_zone_map_filter(const std::vector<const ColumnPredicate*>& predicates,
                                       const ColumnPredicate* del_predicate,

be/src/storage/rowset/column_reader.h

@@ -34,25 +34,21 @@
 
 #pragma once
 
-#include <algorithm>
-#include <bitset>
 #include <cstddef>
 #include <cstdint>
 #include <memory>
 #include <utility>
 
 #include "column/datum.h"
-#include "column/fixed_length_column.h"
-#include "column/vectorized_fwd.h"
 #include "common/statusor.h"
 #include "gen_cpp/segment.pb.h"
-#include "runtime/mem_pool.h"
 #include "storage/index/inverted/inverted_index_iterator.h"
 #include "storage/predicate_tree/predicate_tree_fwd.h"
 #include "storage/range.h"
 #include "storage/rowset/bitmap_index_reader.h"
 #include "storage/rowset/bloom_filter_index_reader.h"
 #include "storage/rowset/common.h"
+#include "storage/rowset/options.h"
 #include "storage/rowset/ordinal_page_index.h"
 #include "storage/rowset/page_handle.h"
 #include "storage/rowset/segment.h"
@@ -154,13 +150,19 @@ class ColumnReader {
 
     int32_t num_data_pages() { return _ordinal_index ? _ordinal_index->num_data_pages() : 0; }
 
-    // page-level zone map filter.
+    // Return the ordinal range of a page
+    std::pair<ordinal_t, ordinal_t> get_page_range(size_t page_index);
 
+    // page-level zone map filter.
     Status zone_map_filter(const std::vector<const ::starrocks::ColumnPredicate*>& p,
                            const ::starrocks::ColumnPredicate* del_predicate,
                            std::unordered_set<uint32_t>* del_partial_filtered_pages, SparseRange<>* row_ranges,
                            const IndexReadOptions& opts, CompoundNodeType pred_relation);
 
+    // NOTE: RAW interface should be used carefully
+    // Return all page-level zonemap
+    StatusOr<std::vector<ZoneMapDetail>> get_raw_zone_map(const IndexReadOptions& opts);
+
     // segment-level zone map filter.
     // Return false to filter out this segment.
     // same as `match_condition`, used by vector engine.