diff --git a/ibis/backends/base/sql/compiler/query_builder.py b/ibis/backends/base/sql/compiler/query_builder.py index cb968f01fef3..ee6f4f76fe13 100644 --- a/ibis/backends/base/sql/compiler/query_builder.py +++ b/ibis/backends/base/sql/compiler/query_builder.py @@ -100,9 +100,11 @@ def _format_table(self, op): ctx = self.context orig_op = op - if isinstance(op, ops.SelfReference): + if isinstance(op, (ops.SelfReference, ops.Sample)): op = op.table + alias = ctx.get_ref(orig_op) + if isinstance(op, ops.InMemoryTable): result = self._format_in_memory_table(op) elif isinstance(op, ops.PhysicalTable): @@ -117,26 +119,28 @@ def _format_table(self, op): db=getattr(op, "namespace", None), quoted=self.parent.translator_class._quote_identifiers, ).sql(dialect=self.parent.translator_class._dialect_name) + elif ctx.is_extracted(op): + if isinstance(orig_op, ops.SelfReference): + result = ctx.get_ref(op) + else: + result = alias else: - # A subquery - if ctx.is_extracted(op): - # Was put elsewhere, e.g. WITH block, we just need to grab its - # alias - alias = ctx.get_ref(orig_op) - - # HACK: self-references have to be treated more carefully here - if isinstance(orig_op, ops.SelfReference): - return f"{ctx.get_ref(op)} {alias}" - else: - return alias - - subquery = ctx.get_compiled_expr(orig_op) + subquery = ctx.get_compiled_expr(op) result = f"(\n{util.indent(subquery, self.indent)}\n)" - result += f" {ctx.get_ref(orig_op)}" + if result != alias: + result = f"{result} {alias}" + + if isinstance(orig_op, ops.Sample): + result = self._format_sample(orig_op, result) return result + def _format_sample(self, op, table): + # Should never be hit in practice, as Sample operations should be rewritten + # before this point for all backends without TABLESAMPLE support + raise com.UnsupportedOperationError("`Table.sample` is not supported") + def get_result(self): # Got to unravel the join stack; the nesting order could be # arbitrary, so we do a depth first search and push the join tokens diff --git a/ibis/backends/impala/compiler.py b/ibis/backends/impala/compiler.py index 897c6a619979..558ab877819b 100644 --- a/ibis/backends/impala/compiler.py +++ b/ibis/backends/impala/compiler.py @@ -3,6 +3,7 @@ import ibis.expr.operations as ops from ibis.backends.base.sql.compiler import Compiler, ExprTranslator, TableSetFormatter from ibis.backends.base.sql.registry import binary_infix_ops, operation_registry, unary +from ibis.expr.rewrites import rewrite_sample class ImpalaTableSetFormatter(TableSetFormatter): @@ -58,3 +59,4 @@ def _floor_divide(op): class ImpalaCompiler(Compiler): translator_class = ImpalaExprTranslator table_set_formatter_class = ImpalaTableSetFormatter + rewrites = Compiler.rewrites | rewrite_sample diff --git a/ibis/backends/tests/test_generic.py b/ibis/backends/tests/test_generic.py index bdc5ae95f5ca..e5133f6a5550 100644 --- a/ibis/backends/tests/test_generic.py +++ b/ibis/backends/tests/test_generic.py @@ -1535,7 +1535,6 @@ def test_dynamic_table_slice_with_computed_offset(backend): "datafusion", "druid", "flink", - "impala", "polars", "snowflake", ] @@ -1561,7 +1560,6 @@ def test_sample(backend): "datafusion", "druid", "flink", - "impala", "polars", "snowflake", ]