From 155cf6e5a576b974deeccd782ded08c7cb1f2c04 Mon Sep 17 00:00:00 2001
From: qianmoQ <shicheng@devlive.org>
Date: Wed, 21 Feb 2024 17:23:53 +0800
Subject: [PATCH] [Core] [DataSet] Support sampling key

---
 .../datacap/service/service/impl/DataSetServiceImpl.java | 1 +
 .../src/main/kotlin/io/edurt/datacap/sql/AbstractSql.kt  | 9 +++++++++
 .../kotlin/io/edurt/datacap/sql/builder/TableBuilder.kt  | 4 ++++
 3 files changed, 14 insertions(+)

diff --git a/core/datacap-service/src/main/java/io/edurt/datacap/service/service/impl/DataSetServiceImpl.java b/core/datacap-service/src/main/java/io/edurt/datacap/service/service/impl/DataSetServiceImpl.java
index c8cf05d2..1b342e8a 100644
--- a/core/datacap-service/src/main/java/io/edurt/datacap/service/service/impl/DataSetServiceImpl.java
+++ b/core/datacap-service/src/main/java/io/edurt/datacap/service/service/impl/DataSetServiceImpl.java
@@ -394,6 +394,7 @@ public class DataSetServiceImpl
             TableBuilder.Companion.ORDER_BY(columnEntities.stream().filter(DataSetColumnEntity::isOrderByKey).map(DataSetColumnEntity::getName).collect(Collectors.toList()));
             TableBuilder.Companion.PARTITION_BY(columnEntities.stream().filter(DataSetColumnEntity::isPartitionKey).map(DataSetColumnEntity::getName).collect(Collectors.toList()));
             TableBuilder.Companion.PRIMARY_KEY(columnEntities.stream().filter(DataSetColumnEntity::isPrimaryKey).map(DataSetColumnEntity::getName).collect(Collectors.toList()));
+            TableBuilder.Companion.SAMPLING_KEY(columnEntities.stream().filter(DataSetColumnEntity::isSamplingKey).map(DataSetColumnEntity::getName).collect(Collectors.toList()));
             String sql = TableBuilder.Companion.SQL();
             log.info("Create table sql \n {} \n on dataset [ {} ]", sql, entity.getName());
 
diff --git a/core/datacap-sql/src/main/kotlin/io/edurt/datacap/sql/AbstractSql.kt b/core/datacap-sql/src/main/kotlin/io/edurt/datacap/sql/AbstractSql.kt
index 17bfbe7b..d8307d2b 100644
--- a/core/datacap-sql/src/main/kotlin/io/edurt/datacap/sql/AbstractSql.kt
+++ b/core/datacap-sql/src/main/kotlin/io/edurt/datacap/sql/AbstractSql.kt
@@ -163,6 +163,11 @@ abstract class AbstractSql<T> {
         return getSelf()
     }
 
+    fun SAMPING_KEY(columns: List<String>): T {
+        sql().samplingKey.addAll(columns.map { item -> "\t$item" })
+        return getSelf()
+    }
+
     fun FROM(table: String?): T {
         sql().tables.add(table)
         return getSelf()
@@ -533,6 +538,7 @@ abstract class AbstractSql<T> {
         var orderByKey: MutableList<String?> = ArrayList()
         val partitionByKey: MutableList<String?> = ArrayList()
         val primaryKey: MutableList<String?> = ArrayList()
+        val samplingKey: MutableList<String?> = ArrayList()
 
         init {
             // Prevent Synthetic Access
@@ -688,6 +694,9 @@ abstract class AbstractSql<T> {
             if (primaryKey.isNotEmpty()) {
                 sqlClause(builder, "PRIMARY KEY", primaryKey, "(", ")", ", ")
             }
+            if (samplingKey.isNotEmpty()) {
+                sqlClause(builder, "SAMPLE BY", samplingKey, "(", ")", ", ")
+            }
             if (end) {
                 builder.append(";")
             }
diff --git a/core/datacap-sql/src/main/kotlin/io/edurt/datacap/sql/builder/TableBuilder.kt b/core/datacap-sql/src/main/kotlin/io/edurt/datacap/sql/builder/TableBuilder.kt
index 65b3982f..caac64bf 100644
--- a/core/datacap-sql/src/main/kotlin/io/edurt/datacap/sql/builder/TableBuilder.kt
+++ b/core/datacap-sql/src/main/kotlin/io/edurt/datacap/sql/builder/TableBuilder.kt
@@ -46,6 +46,10 @@ class TableBuilder {
             sql().PRIMARY_KEY(values)
         }
 
+        fun SAMPLING_KEY(values: List<String>) {
+            sql().SAMPING_KEY(values)
+        }
+
         fun END() {
             sql().END()
         }