[UT] Add regression tests for lake scan sample options

Signed-off-by: levin-kitty <levin.de.kitty@gmail.com>

Author: levin-kitty

be/test/storage/lake/lake_data_source_test.cpp

@@ -35,6 +35,7 @@
 #include "exec/exec_env.h"
 #include "exec/pipeline/fragment_context.h"
 #include "exec/pipeline/scan/morsel.h"
+#include "exec/runtime/fragment_context.h"
 #include "fs/fs_util.h"
 #include "runtime/descriptor_helper.h"
 #include "runtime/descriptors.h"
@@ -49,6 +50,7 @@
 #include "storage/lake/tablet_writer.h"
 #include "storage/lake/update_manager.h"
 #include "storage/rowset/base_rowset.h"
+#include "storage/rowset/segment_options.h"
 #include "storage/storage_env.h"
 #include "storage/tablet_schema.h"
 #include "storage_primitive/vector_search_option.h"
@@ -378,6 +380,109 @@ TEST_F(LakeDataSourceTest, get_tablet_schema) {
     ASSERT_EQ(schema_id, tablet_schema->id());
 }
 
+TEST_F(LakeDataSourceTest, propagate_sample_options) {
+    SyncPoint::GetInstance()->EnableProcessing();
+    DeferOp defer([] {
+        SyncPoint::GetInstance()->ClearAllCallBacks();
+        SyncPoint::GetInstance()->DisableProcessing();
+    });
+
+    create_rowsets_for_testing(_tablet_metadata.get(), 2);
+
+    TTableSampleOptions sample_options;
+    sample_options.__set_enable_sampling(true);
+    sample_options.__set_sample_method(SampleMethod::BY_BLOCK);
+    sample_options.__set_random_seed(7);
+    sample_options.__set_probability_percent_v2(0.5);
+
+    TLakeScanNode lake_scan_node;
+    lake_scan_node.__set_tuple_id(0);
+    lake_scan_node.__set_key_column_name({"c0"});
+    lake_scan_node.__set_key_column_type({TPrimitiveType::INT});
+    lake_scan_node.__set_is_preaggregation(true);
+    lake_scan_node.__set_sample_options(sample_options);
+
+    TPlanNode plan_node;
+    plan_node.__set_node_id(0);
+    plan_node.__set_node_type(TPlanNodeType::LAKE_SCAN_NODE);
+    plan_node.__set_lake_scan_node(lake_scan_node);
+
+    auto runtime_state = create_runtime_state();
+    pipeline::FragmentContext fragment_ctx;
+    runtime_state->set_fragment_ctx(&fragment_ctx, &fragment_ctx.fragment_runtime_state());
+    runtime_state->set_fragment_dict_state(fragment_ctx.dict_state());
+
+    TDescriptorTableBuilder desc_tbl_builder;
+    TSlotDescriptorBuilder slot_desc_builder;
+    auto slot =
+            slot_desc_builder.type(LogicalType::TYPE_INT).column_name("c0").column_pos(0).nullable(true).build();
+
+    TTupleDescriptorBuilder tuple_desc_builder;
+    tuple_desc_builder.add_slot(slot);
+    tuple_desc_builder.build(&desc_tbl_builder);
+
+    DescriptorTbl* desc_tbl = nullptr;
+    ASSERT_OK(DescriptorTbl::create(runtime_state.get(), runtime_state->obj_pool(), desc_tbl_builder.desc_tbl(),
+                                    &desc_tbl, config::vector_chunk_size));
+    runtime_state->set_desc_tbl(desc_tbl);
+
+    TTableDescriptor table_descriptor;
+    table_descriptor.__set_id(0);
+    table_descriptor.__set_tableType(TTableType::OLAP_TABLE);
+    table_descriptor.__set_tableName("test_table");
+    table_descriptor.__set_dbName("test_db");
+
+    auto* table_desc = runtime_state->obj_pool()->add(new OlapTableDescriptor(table_descriptor));
+    desc_tbl->get_tuple_descriptor(0)->set_table_desc(table_desc);
+
+    connector::LakeDataSourceProvider provider(/*scan_node=*/nullptr, plan_node);
+    provider.set_lake_tablet_manager(_tablet_mgr);
+
+    TInternalScanRange internal_scan_range;
+    internal_scan_range.__set_tablet_id(_tablet_metadata->id());
+    internal_scan_range.__set_version(std::to_string(_tablet_metadata->version()));
+
+    TScanRange scan_range;
+    scan_range.__set_internal_scan_range(internal_scan_range);
+
+    ASSIGN_OR_ABORT(auto tablet, _tablet_mgr->get_tablet(_tablet_metadata->id(), _tablet_metadata->version()));
+
+    std::vector<BaseRowsetSharedPtr> base_rowsets;
+    for (auto& rowset : tablet.get_rowsets()) {
+        base_rowsets.emplace_back(rowset);
+    }
+
+    pipeline::ScanMorsel morsel(plan_node.node_id, scan_range);
+    morsel.set_rowsets(base_rowsets);
+
+    bool seen = false;
+    TTableSampleOptions propagated;
+    SyncPoint::GetInstance()->SetCallBack("Rowset::read::seg_options", [&](void* arg) {
+        auto* segment_options = static_cast<SegmentReadOptions*>(arg);
+        propagated = segment_options->sample_options;
+        seen = true;
+    });
+
+    connector::LakeDataSource data_source(&provider, scan_range);
+    RuntimeProfile profile("LakeDataSourceTest");
+    data_source.set_runtime_profile(&profile);
+    data_source.set_morsel(&morsel);
+
+    DeferOp close_guard([&] { data_source.close(runtime_state.get()); });
+    ASSERT_OK(data_source.open(runtime_state.get()));
+
+    ASSERT_TRUE(seen);
+    ASSERT_TRUE(propagated.__isset.enable_sampling);
+    EXPECT_TRUE(propagated.enable_sampling);
+    ASSERT_TRUE(propagated.__isset.sample_method);
+    EXPECT_EQ(SampleMethod::BY_BLOCK, propagated.sample_method);
+    ASSERT_TRUE(propagated.__isset.random_seed);
+    EXPECT_EQ(7, propagated.random_seed);
+    ASSERT_TRUE(propagated.__isset.probability_percent_v2);
+    EXPECT_DOUBLE_EQ(0.5, propagated.probability_percent_v2);
+    EXPECT_FALSE(propagated.__isset.probability_percent);
+}
+
 // Exercise the vector-search branches of LakeDataSource::open() that were added by
 // the shared-data vector index read PR:
 //   * open() propagation of TVectorSearchOptions from thrift into _use_vector_index etc.

fe/fe-core/src/test/java/com/starrocks/sql/plan/LakeTableSamplePlanTest.java

@@ -0,0 +1,75 @@
+// Copyright 2021-present StarRocks, Inc. All rights reserved.
+//
+// Licensed 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
+//
+//     https://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.
+
+package com.starrocks.sql.plan;
+
+import com.starrocks.common.util.UUIDUtil;
+import com.starrocks.planner.OlapScanNode;
+import com.starrocks.planner.ScanNode;
+import com.starrocks.qe.ConnectContext;
+import com.starrocks.server.RunMode;
+import com.starrocks.thrift.SampleMethod;
+import com.starrocks.thrift.TPlanNode;
+import com.starrocks.thrift.TPlanNodeType;
+import com.starrocks.thrift.TTableSampleOptions;
+import com.starrocks.utframe.StarRocksAssert;
+import com.starrocks.utframe.UtFrameUtils;
+import org.junit.jupiter.api.Assertions;
+import org.junit.jupiter.api.BeforeAll;
+import org.junit.jupiter.api.Test;
+
+public class LakeTableSamplePlanTest {
+    private static ConnectContext connectContext;
+
+    @BeforeAll
+    public static void setUp() throws Exception {
+        UtFrameUtils.createMinStarRocksCluster(RunMode.SHARED_DATA);
+        connectContext = UtFrameUtils.createDefaultCtx();
+
+        StarRocksAssert starRocksAssert = new StarRocksAssert(connectContext);
+        starRocksAssert.withDatabase("lake_sample_test")
+                .useDatabase("lake_sample_test")
+                .withTable("CREATE TABLE lake_sample_test.t0 (c0 INT) " +
+                        "DUPLICATE KEY(c0) " +
+                        "DISTRIBUTED BY HASH(c0) BUCKETS 1");
+    }
+
+    @Test
+    public void testLakeScanNodeCarriesSampleOptions() throws Exception {
+        connectContext.setQueryId(UUIDUtil.genUUID());
+        connectContext.setExecutionId(UUIDUtil.toTUniqueId(connectContext.getQueryId()));
+
+        String sql = "SELECT c0 FROM t0 " +
+                "SAMPLE('method'='by_block', 'seed'='7', 'percent'='0.5')";
+        ExecPlan execPlan = UtFrameUtils.getPlanAndFragment(connectContext, sql).second;
+
+        Assertions.assertEquals(1, execPlan.getScanNodes().size());
+        ScanNode scanNode = execPlan.getScanNodes().get(0);
+        Assertions.assertInstanceOf(OlapScanNode.class, scanNode);
+
+        TPlanNode thriftNode = scanNode.treeToThrift().getNodes().get(0);
+        Assertions.assertEquals(TPlanNodeType.LAKE_SCAN_NODE, thriftNode.getNode_type());
+        Assertions.assertTrue(thriftNode.isSetLake_scan_node());
+        Assertions.assertTrue(thriftNode.getLake_scan_node().isSetSample_options());
+
+        TTableSampleOptions sampleOptions =
+                thriftNode.getLake_scan_node().getSample_options();
+        Assertions.assertTrue(sampleOptions.isEnable_sampling());
+        Assertions.assertEquals(SampleMethod.BY_BLOCK, sampleOptions.getSample_method());
+        Assertions.assertEquals(7L, sampleOptions.getRandom_seed());
+        Assertions.assertTrue(sampleOptions.isSetProbability_percent_v2());
+        Assertions.assertEquals(
+                0.5, sampleOptions.getProbability_percent_v2(), 0.000001);
+    }
+}