This is an automated email from the ASF dual-hosted git repository. jhyde pushed a commit to branch main in repository https://gitbox.apache.org/repos/asf/calcite.git
commit 50c3edfc3d6630528ab51fe836bd50df82cc7db8 Author: Julian Hyde <[email protected]> AuthorDate: Sat Aug 12 11:17:19 2023 -0700 [CALCITE-5916] In RelBuilder, add sample() method (equivalent to SQL TABLESAMPLE clause) The sample method optimizes to empty when rate = 0.0, and optimizes to no sampling when rate = 1.0. --- .../org/apache/calcite/rel/core/RelFactories.java | 30 ++++++++++ .../java/org/apache/calcite/rel/core/Sample.java | 5 -- .../apache/calcite/sql2rel/SqlToRelConverter.java | 34 +++-------- .../java/org/apache/calcite/tools/RelBuilder.java | 37 ++++++++++++ .../org/apache/calcite/test/RelBuilderTest.java | 65 ++++++++++++++++++++++ site/_docs/algebra.md | 1 + 6 files changed, 142 insertions(+), 30 deletions(-) diff --git a/core/src/main/java/org/apache/calcite/rel/core/RelFactories.java b/core/src/main/java/org/apache/calcite/rel/core/RelFactories.java index 3a0d543351..a00171ef76 100644 --- a/core/src/main/java/org/apache/calcite/rel/core/RelFactories.java +++ b/core/src/main/java/org/apache/calcite/rel/core/RelFactories.java @@ -20,6 +20,7 @@ import org.apache.calcite.linq4j.function.Experimental; import org.apache.calcite.plan.Context; import org.apache.calcite.plan.Contexts; import org.apache.calcite.plan.RelOptCluster; +import org.apache.calcite.plan.RelOptSamplingParameters; import org.apache.calcite.plan.RelOptTable; import org.apache.calcite.plan.RelTraitSet; import org.apache.calcite.rel.RelCollation; @@ -99,6 +100,9 @@ public class RelFactories { public static final AggregateFactory DEFAULT_AGGREGATE_FACTORY = new AggregateFactoryImpl(); + public static final SampleFactory DEFAULT_SAMPLE_FACTORY = + new SampleFactoryImpl(); + public static final MatchFactory DEFAULT_MATCH_FACTORY = new MatchFactoryImpl(); @@ -137,6 +141,7 @@ public class RelFactories { DEFAULT_TABLE_SCAN_FACTORY, DEFAULT_TABLE_FUNCTION_SCAN_FACTORY, DEFAULT_SNAPSHOT_FACTORY, + DEFAULT_SAMPLE_FACTORY, DEFAULT_MATCH_FACTORY, DEFAULT_SPOOL_FACTORY, DEFAULT_REPEAT_UNION_FACTORY); @@ -601,6 +606,26 @@ public class RelFactories { } } + /** + * Can create a {@link Sample} of + * the appropriate type for a rule's calling convention. + */ + public interface SampleFactory { + /** Creates a {@link Sample}. */ + RelNode createSample(RelNode input, RelOptSamplingParameters parameter); + } + + /** + * Implementation of {@link SampleFactory} + * that returns a {@link Sample}. + */ + private static class SampleFactoryImpl implements SampleFactory { + @Override public RelNode createSample(RelNode input, + RelOptSamplingParameters parameter) { + return new Sample(input.getCluster(), input, parameter); + } + } + /** * Can create a {@link Spool} of * the appropriate type for a rule's calling convention. @@ -661,6 +686,7 @@ public class RelFactories { public final TableFunctionScanFactory tableFunctionScanFactory; public final SnapshotFactory snapshotFactory; public final MatchFactory matchFactory; + public final SampleFactory sampleFactory; public final SpoolFactory spoolFactory; public final RepeatUnionFactory repeatUnionFactory; @@ -677,6 +703,7 @@ public class RelFactories { TableScanFactory scanFactory, TableFunctionScanFactory tableFunctionScanFactory, SnapshotFactory snapshotFactory, + SampleFactory sampleFactory, MatchFactory matchFactory, SpoolFactory spoolFactory, RepeatUnionFactory repeatUnionFactory) { @@ -694,6 +721,7 @@ public class RelFactories { this.tableFunctionScanFactory = requireNonNull(tableFunctionScanFactory, "tableFunctionScanFactory"); this.snapshotFactory = requireNonNull(snapshotFactory, "snapshotFactory"); + this.sampleFactory = requireNonNull(sampleFactory, "sampleFactory"); this.matchFactory = requireNonNull(matchFactory, "matchFactory"); this.spoolFactory = requireNonNull(spoolFactory, "spoolFactory"); this.repeatUnionFactory = requireNonNull(repeatUnionFactory, "repeatUnionFactory"); @@ -731,6 +759,8 @@ public class RelFactories { .orElse(DEFAULT_TABLE_FUNCTION_SCAN_FACTORY), context.maybeUnwrap(SnapshotFactory.class) .orElse(DEFAULT_SNAPSHOT_FACTORY), + context.maybeUnwrap(SampleFactory.class) + .orElse(DEFAULT_SAMPLE_FACTORY), context.maybeUnwrap(MatchFactory.class) .orElse(DEFAULT_MATCH_FACTORY), context.maybeUnwrap(SpoolFactory.class) diff --git a/core/src/main/java/org/apache/calcite/rel/core/Sample.java b/core/src/main/java/org/apache/calcite/rel/core/Sample.java index 112291e5bb..a4dbfacea1 100644 --- a/core/src/main/java/org/apache/calcite/rel/core/Sample.java +++ b/core/src/main/java/org/apache/calcite/rel/core/Sample.java @@ -25,13 +25,8 @@ import org.apache.calcite.rel.RelNode; import org.apache.calcite.rel.RelWriter; import org.apache.calcite.rel.SingleRel; -import org.checkerframework.checker.nullness.qual.Nullable; - import java.math.BigDecimal; import java.util.List; -import java.util.Objects; - -import static java.util.Objects.requireNonNull; /** * Relational expression that returns a sample of the rows from its input. diff --git a/core/src/main/java/org/apache/calcite/sql2rel/SqlToRelConverter.java b/core/src/main/java/org/apache/calcite/sql2rel/SqlToRelConverter.java index 09a5a5479f..45853db7e3 100644 --- a/core/src/main/java/org/apache/calcite/sql2rel/SqlToRelConverter.java +++ b/core/src/main/java/org/apache/calcite/sql2rel/SqlToRelConverter.java @@ -23,7 +23,6 @@ import org.apache.calcite.linq4j.Ord; import org.apache.calcite.linq4j.tree.TableExpressionFactory; import org.apache.calcite.plan.RelOptCluster; import org.apache.calcite.plan.RelOptPlanner; -import org.apache.calcite.plan.RelOptSamplingParameters; import org.apache.calcite.plan.RelOptTable; import org.apache.calcite.plan.RelOptUtil; import org.apache.calcite.plan.RelTraitSet; @@ -50,7 +49,6 @@ import org.apache.calcite.rel.core.JoinInfo; import org.apache.calcite.rel.core.JoinRelType; import org.apache.calcite.rel.core.Project; import org.apache.calcite.rel.core.RelFactories; -import org.apache.calcite.rel.core.Sample; import org.apache.calcite.rel.core.Sort; import org.apache.calcite.rel.hint.HintStrategyTable; import org.apache.calcite.rel.hint.Hintable; @@ -2362,29 +2360,15 @@ public class SqlToRelConverter { (SqlSampleSpec.SqlTableSampleSpec) sampleSpec; convertFrom(bb, operands.get(0)); - // Treat TABLESAMPLE(0) and TABLESAMPLE(100) as no table - // sampling at all. Not strictly correct: TABLESAMPLE(0) - // should produce no output, but it simplifies implementation - // to know that some amount of sampling will occur. - // In practice values less than ~1E-43% are treated as 0.0 and - // values greater than ~99.999997% are treated as 1.0 - relBuilder.push(bb.root()); - if (tableSampleSpec.sampleRate.compareTo(BigDecimal.ZERO) == 0) { - // The table sample rate is 0; the query should return empty. - relBuilder.empty(); - } else if (tableSampleSpec.sampleRate.compareTo(BigDecimal.ONE) == 0) { - // The table sample rate is 1; the query should return the contents - // of the underlying table. - } else { - RelOptSamplingParameters params = - new RelOptSamplingParameters( - tableSampleSpec.isBernoulli(), - tableSampleSpec.sampleRate, - tableSampleSpec.isRepeatable(), - tableSampleSpec.getRepeatableSeed()); - relBuilder.push(new Sample(cluster, relBuilder.build(), params)); - } - bb.setRoot(relBuilder.build(), true); + bb.setRoot( + relBuilder.push(bb.root()) + .sample(tableSampleSpec.isBernoulli(), + tableSampleSpec.sampleRate, + tableSampleSpec.isRepeatable() + ? tableSampleSpec.getRepeatableSeed() + : null) + .build(), + true); } else { throw new AssertionError("unknown TABLESAMPLE type: " + sampleSpec); } diff --git a/core/src/main/java/org/apache/calcite/tools/RelBuilder.java b/core/src/main/java/org/apache/calcite/tools/RelBuilder.java index 1c07a10c3f..3e45b0cc92 100644 --- a/core/src/main/java/org/apache/calcite/tools/RelBuilder.java +++ b/core/src/main/java/org/apache/calcite/tools/RelBuilder.java @@ -23,6 +23,7 @@ import org.apache.calcite.plan.Contexts; import org.apache.calcite.plan.Convention; import org.apache.calcite.plan.RelOptCluster; import org.apache.calcite.plan.RelOptPredicateList; +import org.apache.calcite.plan.RelOptSamplingParameters; import org.apache.calcite.plan.RelOptSchema; import org.apache.calcite.plan.RelOptTable; import org.apache.calcite.plan.RelOptUtil; @@ -47,6 +48,7 @@ import org.apache.calcite.rel.core.Minus; import org.apache.calcite.rel.core.Project; import org.apache.calcite.rel.core.RelFactories; import org.apache.calcite.rel.core.RepeatUnion; +import org.apache.calcite.rel.core.Sample; import org.apache.calcite.rel.core.Snapshot; import org.apache.calcite.rel.core.Sort; import org.apache.calcite.rel.core.Spool; @@ -3560,6 +3562,41 @@ public class RelBuilder { return project(exprList); } + /** Creates a {@link Sample}. (Repeatable if seed is not null.) */ + public RelBuilder sample(boolean bernoulli, BigDecimal sampleRate, + @Nullable Integer repeatableSeed) { + boolean repeatable; + int seed; + if (repeatableSeed != null) { + repeatable = true; + seed = repeatableSeed; + } else { + repeatable = false; + seed = 0; + } + return sample(bernoulli, sampleRate, repeatable, seed); + } + + /** Creates a {@link Sample}. */ + private RelBuilder sample(boolean bernoulli, BigDecimal sampleRate, + boolean repeatable, int repeatableSeed) { + if (sampleRate.compareTo(BigDecimal.ZERO) == 0) { + // The sample rate is 0%; the query should return empty. + return empty(); + } else if (sampleRate.compareTo(BigDecimal.ONE) == 0) { + // The table sample rate is 100%; the query should return the contents + // of the underlying table. + return this; + } else { + final Frame frame = stack.pop(); + final RelNode r = frame.rel; + final RelOptSamplingParameters param = + new RelOptSamplingParameters(bernoulli, sampleRate, repeatable, + repeatableSeed); + return push(struct.sampleFactory.createSample(r, param)); + } + } + /** Creates a {@link Match}. */ public RelBuilder match(RexNode pattern, boolean strictStart, boolean strictEnd, Map<String, RexNode> patternDefinitions, diff --git a/core/src/test/java/org/apache/calcite/test/RelBuilderTest.java b/core/src/test/java/org/apache/calcite/test/RelBuilderTest.java index d24b7a2674..901d37dc2e 100644 --- a/core/src/test/java/org/apache/calcite/test/RelBuilderTest.java +++ b/core/src/test/java/org/apache/calcite/test/RelBuilderTest.java @@ -100,6 +100,7 @@ import org.junit.jupiter.params.ParameterizedTest; import org.junit.jupiter.params.provider.ValueSource; import java.lang.reflect.Method; +import java.math.BigDecimal; import java.sql.Connection; import java.sql.DriverManager; import java.sql.PreparedStatement; @@ -3905,6 +3906,70 @@ public class RelBuilderTest { assertThat(f.apply(createBuilder(), false), hasTree(expectedExcludeNulls)); } + @Test void testSample() { + // Equivalent SQL: + // SELECT * + // FROM emp + // TABLESAMPLE SYSTEM(40) + final Function<RelBuilder, RelNode> f = + b -> b.scan("EMP") + .sample(false, new BigDecimal("0.4"), null) + .build(); + final String expected = "" + + "Sample(mode=[system], rate=[0.4], repeatableSeed=[-])\n" + + " LogicalTableScan(table=[[scott, EMP]])\n"; + assertThat(f.apply(createBuilder()), hasTree(expected)); + } + + @Test void testSampleBernoulliRepeatable() { + // Equivalent SQL: + // SELECT * + // FROM emp + // TABLESAMPLE BERNOULLI(25, 31415926) + final Function<RelBuilder, RelNode> f = + b -> b.scan("EMP") + .sample(true, new BigDecimal("0.25"), 31_415_926) + .build(); + final String expected = "" + + "Sample(mode=[bernoulli], rate=[0.25], repeatableSeed=[31415926])\n" + + " LogicalTableScan(table=[[scott, EMP]])\n"; + assertThat(f.apply(createBuilder()), hasTree(expected)); + } + + /** Tests that TABLESAMPLE(0) returns zero rows. */ + @Test void testSampleZero() { + // Equivalent SQL: + // SELECT * + // FROM emp + // TABLESAMPLE SYSTEM(0) + final BiFunction<RelBuilder, Boolean, RelNode> f = + (b, mode) -> b.scan("EMP") + .sample(mode, BigDecimal.ZERO, null) + .build(); + final String expected = "LogicalValues(tuples=[[]])\n"; + assertThat(f.apply(createBuilder(), true), hasTree(expected)); + assertThat(f.apply(createBuilder(), false), hasTree(expected)); + } + + /** Tests that TABLESAMPLE(100) (rate=1.0) does no sampling. */ + @Test void testSampleAll() { + // Equivalent SQL: + // SELECT * + // FROM emp + // TABLESAMPLE SYSTEM(100) + // becomes + // SELECT * + // FROM emp + final BiFunction<RelBuilder, Boolean, RelNode> f = + (b, mode) -> b.scan("EMP") + .sample(mode, BigDecimal.ONE, null) + .build(); + final String expected = "" + + "LogicalTableScan(table=[[scott, EMP]])\n"; + assertThat(f.apply(createBuilder(), true), hasTree(expected)); + assertThat(f.apply(createBuilder(), false), hasTree(expected)); + } + @Test void testMatchRecognize() { // Equivalent SQL: // SELECT * diff --git a/site/_docs/algebra.md b/site/_docs/algebra.md index 552a139d8a..9d1899e563 100644 --- a/site/_docs/algebra.md +++ b/site/_docs/algebra.md @@ -351,6 +351,7 @@ return the `RelBuilder`. | `intersect(all [, n])` | Creates an [Intersect]({{ site.apiRoot }}/org/apache/calcite/rel/core/Intersect.html) of the `n` (default two) most recent relational expressions. | `minus(all)` | Creates a [Minus]({{ site.apiRoot }}/org/apache/calcite/rel/core/Minus.html) of the two most recent relational expressions. | `repeatUnion(tableName, all [, n])` | Creates a [RepeatUnion]({{ site.apiRoot }}/org/apache/calcite/rel/core/RepeatUnion.html) associated to a [TransientTable]({{ site.apiRoot }}/org/apache/calcite/schema/TransientTable.html) of the two most recent relational expressions, with `n` maximum number of iterations (default -1, i.e. no limit). +| `sample(bernoulli, rate [, repeatableSeed])` | Creates a [sample]({{ site.apiRoot }}/org/apache/calcite/rel/core/Sample.html) of at given sampling rate. | `snapshot(period)` | Creates a [Snapshot]({{ site.apiRoot }}/org/apache/calcite/rel/core/Snapshot.html) of the given snapshot period. | `match(pattern, strictStart,` `strictEnd, patterns, measures,` `after, subsets, allRows,` `partitionKeys, orderKeys,` `interval)` | Creates a [Match]({{ site.apiRoot }}/org/apache/calcite/rel/core/Match.html).
