[sql] generate basic TABLESAMPLE

databrickslabs · Aug 21, 2024 · 06911f4 · 06911f4
1 parent b08f628
commit 06911f4
Show file tree

Hide file tree

Showing 2 changed files with 21 additions and 0 deletions.
diff --git a/core/src/main/scala/com/databricks/labs/remorph/generators/sql/LogicalPlanGenerator.scala b/core/src/main/scala/com/databricks/labs/remorph/generators/sql/LogicalPlanGenerator.scala
@@ -1,6 +1,7 @@
 package com.databricks.labs.remorph.generators.sql
 
 import com.databricks.labs.remorph.generators.{Generator, GeneratorContext}
+import com.databricks.labs.remorph.parsers.intermediate.{BlockSampling, RowSamplingFixedAmount, RowSamplingProbabilistic}
 import com.databricks.labs.remorph.parsers.{intermediate => ir}
 import com.databricks.labs.remorph.transpilers.TranspileException
 
@@ -33,11 +34,23 @@ class LogicalPlanGenerator(val expr: ExpressionGenerator, val explicitDistinct:
     case i: ir.InsertIntoTable => insert(ctx, i)
     case ir.DeleteFromTable(target, None, where, None, None) => delete(ctx, target, where)
     case c: ir.CreateTableCommand => createTable(ctx, c)
+    case t: ir.TableSample => tableSample(ctx, t)
     case ir.NoopNode => ""
     case null => "" // don't fail transpilation if the plan is null
     case x => throw unknown(x)
   }
 
+  // @see https://docs.databricks.com/en/sql/language-manual/sql-ref-syntax-qry-select-sampling.html
+  private def tableSample(ctx: GeneratorContext, t: ir.TableSample): String = {
+    val sampling = t.samplingMethod match {
+      case RowSamplingProbabilistic(probability) => s"$probability PERCENT"
+      case RowSamplingFixedAmount(amount) => s"$amount ROWS"
+      case BlockSampling(probability) => s"BUCKET $probability OUT OF 1"
+    }
+    val seed = t.seed.map(s => s" REPEATABLE ($s)").getOrElse("")
+    s"(${generate(ctx, t.child)}) TABLESAMPLE ($sampling)$seed"
+  }
+
   private def createTable(ctx: GeneratorContext, createTable: ir.CreateTableCommand): String = {
     val columns = createTable.columns
       .map { col =>

diff --git a/.../src/test/scala/com/databricks/labs/remorph/generators/sql/LogicalPlanGeneratorTest.scala b/.../src/test/scala/com/databricks/labs/remorph/generators/sql/LogicalPlanGeneratorTest.scala
@@ -360,4 +360,12 @@ class LogicalPlanGeneratorTest extends AnyWordSpec with GeneratorTestCommon[ir.L
     }
   }
 
+  "TableSample" should {
+    "transpile to TABLESAMPLE" in {
+      ir.TableSample(
+        namedTable("t1"),
+        ir.RowSamplingFixedAmount(BigDecimal(10)),
+        Some(BigDecimal(10))) generates "(t1) TABLESAMPLE (10 ROWS) REPEATABLE (10)"
+    }
+  }
 }