feat(bigquery, impala, mssql, oracle, postgres): compile `Table.sampl…

…e` to native `TABLESAMPLE` syntax when possible

ibis/backends/sql/compilers/bigquery/__init__.py

@@ -22,6 +22,7 @@
     exclude_unsupported_window_frame_from_ops,
     exclude_unsupported_window_frame_from_rank,
     exclude_unsupported_window_frame_from_row_number,
+    lower_sample,
     split_select_distinct_with_order_by,
 )
 from ibis.common.temporal import DateUnit, IntervalUnit, TimestampUnit, TimeUnit
@@ -118,6 +119,14 @@ class BigQueryCompiler(SQLGlotCompiler):
 
     supports_qualify = True
 
+    LOWERED_OPS = {
+        ops.Sample: lower_sample(
+            supports_methods=("block",),
+            supports_seed=False,
+            physical_tables_only=True,
+        ),
+    }
+
     UNSUPPORTED_OPS = (
         ops.DateDiff,
         ops.ExtractAuthority,

ibis/backends/sql/compilers/druid.py

@@ -65,6 +65,7 @@ class DruidCompiler(SQLGlotCompiler):
         ops.TypeOf,
         ops.Unnest,
         ops.Variance,
+        ops.Sample,
     )
 
     SIMPLE_OPS = {

ibis/backends/sql/compilers/impala.py

@@ -10,7 +10,7 @@
 from ibis.backends.sql.compilers.base import NULL, STAR, SQLGlotCompiler
 from ibis.backends.sql.datatypes import ImpalaType
 from ibis.backends.sql.dialects import Impala
-from ibis.backends.sql.rewrites import rewrite_empty_order_by_window
+from ibis.backends.sql.rewrites import lower_sample, rewrite_empty_order_by_window
 
 
 class ImpalaCompiler(SQLGlotCompiler):
@@ -23,6 +23,12 @@ class ImpalaCompiler(SQLGlotCompiler):
         *SQLGlotCompiler.rewrites,
     )
 
+    LOWERED_OPS = {
+        ops.Sample: lower_sample(
+            supports_methods=("block",), physical_tables_only=True
+        ),
+    }
+
     UNSUPPORTED_OPS = (
         ops.ArgMax,
         ops.ArgMin,

ibis/backends/sql/compilers/mssql.py

@@ -22,6 +22,7 @@
 from ibis.backends.sql.rewrites import (
     exclude_unsupported_window_frame_from_ops,
     exclude_unsupported_window_frame_from_row_number,
+    lower_sample,
     p,
     replace,
     split_select_distinct_with_order_by,
@@ -73,6 +74,12 @@ class MSSQLCompiler(SQLGlotCompiler):
     post_rewrites = (split_select_distinct_with_order_by,)
     copy_func_args = True
 
+    LOWERED_OPS = {
+        ops.Sample: lower_sample(
+            supports_methods=("block",), physical_tables_only=True
+        ),
+    }
+
     UNSUPPORTED_OPS = (
         ops.ApproxMedian,
         ops.ArgMax,

ibis/backends/sql/compilers/oracle.py

@@ -16,6 +16,7 @@
     exclude_unsupported_window_frame_from_row_number,
     lower_log2,
     lower_log10,
+    lower_sample,
     rewrite_empty_order_by_window,
 )
 
@@ -46,6 +47,7 @@ class OracleCompiler(SQLGlotCompiler):
     LOWERED_OPS = {
         ops.Log2: lower_log2,
         ops.Log10: lower_log10,
+        ops.Sample: lower_sample(physical_tables_only=True),
     }
 
     UNSUPPORTED_OPS = (

ibis/backends/sql/compilers/postgres.py

@@ -17,7 +17,7 @@
 from ibis.backends.sql.compilers.base import NULL, STAR, AggGen, SQLGlotCompiler
 from ibis.backends.sql.datatypes import PostgresType
 from ibis.backends.sql.dialects import Postgres
-from ibis.backends.sql.rewrites import split_select_distinct_with_order_by
+from ibis.backends.sql.rewrites import lower_sample, split_select_distinct_with_order_by
 from ibis.common.exceptions import InvalidDecoratorError
 from ibis.util import gen_name
 
@@ -50,6 +50,8 @@ class PostgresCompiler(SQLGlotCompiler):
     POS_INF = sge.Literal.number("'Inf'::double precision")
     NEG_INF = sge.Literal.number("'-Inf'::double precision")
 
+    LOWERED_OPS = {ops.Sample: lower_sample(physical_tables_only=True)}
+
     UNSUPPORTED_OPS = (
         ops.RowID,
         ops.TimeDelta,

ibis/backends/sql/compilers/risingwave.py

@@ -26,6 +26,7 @@ class RisingWaveCompiler(PostgresCompiler):
         ops.RandomUUID,
         ops.MultiQuantile,
         ops.ApproxMultiQuantile,
+        ops.Sample,
         *(
             op
             for op in ALL_OPERATIONS

ibis/backends/sql/dialects.py

@@ -307,18 +307,30 @@ class Tokenizer(Hive.Tokenizer):
         STRING_ESCAPES = ["'"]
 
 
+def tablesample_percent_to_int(self, expr):
+    """Impala's TABLESAMPLE only supports integer percentages."""
+    expr = expr.copy()
+    expr.args["percent"] = sge.convert(round(float(expr.args["percent"].this)))
+    return self.tablesample_sql(expr)
+
+
 class Impala(Hive):
     NULL_ORDERING = "nulls_are_large"
     REGEXP_EXTRACT_DEFAULT_GROUP = 0
+    TABLESAMPLE_SIZE_IS_PERCENT = True
+    ALIAS_POST_TABLESAMPLE = False
 
     class Generator(Hive.Generator):
+        TABLESAMPLE_WITH_METHOD = True
+
         TRANSFORMS = Hive.Generator.TRANSFORMS.copy() | {
             sge.ApproxDistinct: rename_func("ndv"),
             sge.IsNan: rename_func("is_nan"),
             sge.IsInf: rename_func("is_inf"),
             sge.DayOfWeek: rename_func("dayofweek"),
             sge.Interval: lambda self, e: _interval(self, e, quote_arg=False),
             sge.CurrentDate: rename_func("current_date"),
+            sge.TableSample: tablesample_percent_to_int,
         }
 
 

ibis/backends/tests/snapshots/test_sql/test_sample/bigquery-False/block.sql

@@ -0,0 +1,3 @@
+SELECT
+  *
+FROM `test` AS `t0` TABLESAMPLE system (50.0 PERCENT)

ibis/backends/tests/snapshots/test_sql/test_sample/bigquery-False/row.sql

@@ -0,0 +1,5 @@
+SELECT
+  *
+FROM `test` AS `t0`
+WHERE
+  RAND() <= 0.5

ibis/backends/tests/snapshots/test_sql/test_sample/bigquery-True/block.sql

@@ -0,0 +1,11 @@
+SELECT
+  *
+FROM (
+  SELECT
+    *
+  FROM `test` AS `t0`
+  WHERE
+    `t0`.`x` > 10
+) AS `t1`
+WHERE
+  RAND() <= 0.5

ibis/backends/tests/snapshots/test_sql/test_sample/bigquery-True/row.sql

@@ -0,0 +1,11 @@
+SELECT
+  *
+FROM (
+  SELECT
+    *
+  FROM `test` AS `t0`
+  WHERE
+    `t0`.`x` > 10
+) AS `t1`
+WHERE
+  RAND() <= 0.5

ibis/backends/tests/snapshots/test_sql/test_sample/clickhouse-False/block.sql

@@ -0,0 +1,5 @@
+SELECT
+  *
+FROM "test" AS "t0"
+WHERE
+  randCanonical() <= 0.5

ibis/backends/tests/snapshots/test_sql/test_sample/clickhouse-False/row.sql

@@ -0,0 +1,5 @@
+SELECT
+  *
+FROM "test" AS "t0"
+WHERE
+  randCanonical() <= 0.5

ibis/backends/tests/snapshots/test_sql/test_sample/clickhouse-True/block.sql

@@ -0,0 +1,11 @@
+SELECT
+  *
+FROM (
+  SELECT
+    *
+  FROM "test" AS "t0"
+  WHERE
+    "t0"."x" > 10
+) AS "t1"
+WHERE
+  randCanonical() <= 0.5

ibis/backends/tests/snapshots/test_sql/test_sample/clickhouse-True/row.sql

@@ -0,0 +1,11 @@
+SELECT
+  *
+FROM (
+  SELECT
+    *
+  FROM "test" AS "t0"
+  WHERE
+    "t0"."x" > 10
+) AS "t1"
+WHERE
+  randCanonical() <= 0.5

ibis/backends/tests/snapshots/test_sql/test_sample/datafusion-False/block.sql

@@ -0,0 +1,5 @@
+SELECT
+  *
+FROM "test" AS "t0"
+WHERE
+  RANDOM() <= 0.5

ibis/backends/tests/snapshots/test_sql/test_sample/datafusion-False/row.sql

@@ -0,0 +1,5 @@
+SELECT
+  *
+FROM "test" AS "t0"
+WHERE
+  RANDOM() <= 0.5

ibis/backends/tests/snapshots/test_sql/test_sample/datafusion-True/block.sql

@@ -0,0 +1,11 @@
+SELECT
+  *
+FROM (
+  SELECT
+    *
+  FROM "test" AS "t0"
+  WHERE
+    "t0"."x" > 10
+) AS "t1"
+WHERE
+  RANDOM() <= 0.5

ibis/backends/tests/snapshots/test_sql/test_sample/datafusion-True/row.sql

@@ -0,0 +1,11 @@
+SELECT
+  *
+FROM (
+  SELECT
+    *
+  FROM "test" AS "t0"
+  WHERE
+    "t0"."x" > 10
+) AS "t1"
+WHERE
+  RANDOM() <= 0.5

ibis/backends/tests/snapshots/test_sql/test_sample/duckdb-False/block.sql

@@ -0,0 +1,3 @@
+SELECT
+  *
+FROM "test" AS "t0" TABLESAMPLE system (50.0 PERCENT)

ibis/backends/tests/snapshots/test_sql/test_sample/duckdb-False/row.sql

@@ -0,0 +1,3 @@
+SELECT
+  *
+FROM "test" AS "t0" TABLESAMPLE bernoulli (50.0 PERCENT)

ibis/backends/tests/snapshots/test_sql/test_sample/duckdb-True/block.sql

@@ -0,0 +1,9 @@
+SELECT
+  *
+FROM (
+  SELECT
+    *
+  FROM "test" AS "t0"
+  WHERE
+    "t0"."x" > 10
+) AS "t1" TABLESAMPLE system (50.0 PERCENT)

ibis/backends/tests/snapshots/test_sql/test_sample/duckdb-True/row.sql

@@ -0,0 +1,9 @@
+SELECT
+  *
+FROM (
+  SELECT
+    *
+  FROM "test" AS "t0"
+  WHERE
+    "t0"."x" > 10
+) AS "t1" TABLESAMPLE bernoulli (50.0 PERCENT)

ibis/backends/tests/snapshots/test_sql/test_sample/exasol-False/block.sql

@@ -0,0 +1,5 @@
+SELECT
+  *
+FROM "test" AS "t0"
+WHERE
+  RANDOM() <= 0.5

ibis/backends/tests/snapshots/test_sql/test_sample/exasol-False/row.sql

@@ -0,0 +1,5 @@
+SELECT
+  *
+FROM "test" AS "t0"
+WHERE
+  RANDOM() <= 0.5

ibis/backends/tests/snapshots/test_sql/test_sample/exasol-True/block.sql

@@ -0,0 +1,11 @@
+SELECT
+  *
+FROM (
+  SELECT
+    *
+  FROM "test" AS "t0"
+  WHERE
+    "t0"."x" > 10
+) AS "t1"
+WHERE
+  RANDOM() <= 0.5

ibis/backends/tests/snapshots/test_sql/test_sample/exasol-True/row.sql

@@ -0,0 +1,11 @@
+SELECT
+  *
+FROM (
+  SELECT
+    *
+  FROM "test" AS "t0"
+  WHERE
+    "t0"."x" > 10
+) AS "t1"
+WHERE
+  RANDOM() <= 0.5

ibis/backends/tests/snapshots/test_sql/test_sample/flink-False/block.sql

@@ -0,0 +1,5 @@
+SELECT
+  *
+FROM `test` AS `t0`
+WHERE
+  RAND() <= 0.5

ibis/backends/tests/snapshots/test_sql/test_sample/flink-False/row.sql

@@ -0,0 +1,5 @@
+SELECT
+  *
+FROM `test` AS `t0`
+WHERE
+  RAND() <= 0.5

ibis/backends/tests/snapshots/test_sql/test_sample/flink-True/block.sql

@@ -0,0 +1,11 @@
+SELECT
+  *
+FROM (
+  SELECT
+    *
+  FROM `test` AS `t0`
+  WHERE
+    `t0`.`x` > 10
+) AS `t1`
+WHERE
+  RAND() <= 0.5

ibis/backends/tests/snapshots/test_sql/test_sample/flink-True/row.sql

@@ -0,0 +1,11 @@
+SELECT
+  *
+FROM (
+  SELECT
+    *
+  FROM `test` AS `t0`
+  WHERE
+    `t0`.`x` > 10
+) AS `t1`
+WHERE
+  RAND() <= 0.5

ibis/backends/tests/snapshots/test_sql/test_sample/impala-False/block.sql

@@ -0,0 +1,3 @@
+SELECT
+  *
+FROM `test` AS `t0` TABLESAMPLE system (50)

ibis/backends/tests/snapshots/test_sql/test_sample/impala-False/row.sql

@@ -0,0 +1,5 @@
+SELECT
+  *
+FROM `test` AS `t0`
+WHERE
+  RAND(UTC_TO_UNIX_MICROS(UTC_TIMESTAMP())) <= 0.5