diff --git a/.github/trigger_files/beam_PostCommit_Java_Tpcds_Dataflow.json b/.github/trigger_files/beam_PostCommit_Java_Tpcds_Dataflow.json new file mode 100644 index 000000000000..0967ef424bce --- /dev/null +++ b/.github/trigger_files/beam_PostCommit_Java_Tpcds_Dataflow.json @@ -0,0 +1 @@ +{} diff --git a/.github/trigger_files/beam_PostCommit_Java_Tpcds_Flink.json b/.github/trigger_files/beam_PostCommit_Java_Tpcds_Flink.json index e3d6056a5de9..b26833333238 100644 --- a/.github/trigger_files/beam_PostCommit_Java_Tpcds_Flink.json +++ b/.github/trigger_files/beam_PostCommit_Java_Tpcds_Flink.json @@ -1,4 +1,4 @@ { "comment": "Modify this file in a trivial way to cause this test suite to run", - "modification": 1 + "modification": 2 } diff --git a/.github/trigger_files/beam_PostCommit_Java_Tpcds_Spark.json b/.github/trigger_files/beam_PostCommit_Java_Tpcds_Spark.json index e3d6056a5de9..b26833333238 100644 --- a/.github/trigger_files/beam_PostCommit_Java_Tpcds_Spark.json +++ b/.github/trigger_files/beam_PostCommit_Java_Tpcds_Spark.json @@ -1,4 +1,4 @@ { "comment": "Modify this file in a trivial way to cause this test suite to run", - "modification": 1 + "modification": 2 } diff --git a/.github/workflows/README.md b/.github/workflows/README.md index c41191e4fe35..178fb6502680 100644 --- a/.github/workflows/README.md +++ b/.github/workflows/README.md @@ -372,6 +372,7 @@ PostCommit Jobs run in a schedule against master branch and generally do not get | [ PostCommit Python Examples Spark ](https://github.com/apache/beam/actions/workflows/beam_PostCommit_Python_Examples_Spark.yml) | ['3.9','3.11','3.12'] |`beam_PostCommit_Python_Examples_Spark.json`| [![.github/workflows/beam_PostCommit_Python_Examples_Spark.yml](https://github.com/apache/beam/actions/workflows/beam_PostCommit_Python_Examples_Spark.yml/badge.svg?event=schedule)](https://github.com/apache/beam/actions/workflows/beam_PostCommit_Python_Examples_Spark.yml?query=event%3Aschedule) | | [ PostCommit Python MongoDBIO IT ](https://github.com/apache/beam/actions/workflows/beam_PostCommit_Python_MongoDBIO_IT.yml) | N/A |`beam_PostCommit_Python_MongoDBIO_IT.json`| [![.github/workflows/beam_PostCommit_Python_MongoDBIO_IT.yml](https://github.com/apache/beam/actions/workflows/beam_PostCommit_Python_MongoDBIO_IT.yml/badge.svg?event=schedule)](https://github.com/apache/beam/actions/workflows/beam_PostCommit_Python_MongoDBIO_IT.yml?query=event%3Aschedule) | | [ PostCommit Python Nexmark Direct ](https://github.com/apache/beam/actions/workflows/beam_PostCommit_Python_Nexmark_Direct.yml) | N/A |`beam_PostCommit_Python_Nexmark_Direct.json`| [![.github/workflows/beam_PostCommit_Python_Nexmark_Direct.yml](https://github.com/apache/beam/actions/workflows/beam_PostCommit_Python_Nexmark_Direct.yml/badge.svg?event=schedule)](https://github.com/apache/beam/actions/workflows/beam_PostCommit_Python_Nexmark_Direct.yml?query=event%3Aschedule) | +| [ PostCommit Python Tpcds Prism ](https://github.com/apache/beam/actions/workflows/beam_PostCommit_Python_Tpcds_Prism.yml) | N/A |`beam_PostCommit_Python_Tpcds_Prism.json`| [![.github/workflows/beam_PostCommit_Python_Tpcds_Prism.yml](https://github.com/apache/beam/actions/workflows/beam_PostCommit_Python_Tpcds_Prism.yml/badge.svg?event=schedule)](https://github.com/apache/beam/actions/workflows/beam_PostCommit_Python_Tpcds_Prism.yml?query=event%3Aschedule) | | [ PostCommit Python ValidatesContainer Dataflow ](https://github.com/apache/beam/actions/workflows/beam_PostCommit_Python_ValidatesContainer_Dataflow.yml) | ['3.9','3.10','3.11','3.12'] |`beam_PostCommit_Python_ValidatesContainer_Dataflow.json`| [![.github/workflows/beam_PostCommit_Python_ValidatesContainer_Dataflow.yml](https://github.com/apache/beam/actions/workflows/beam_PostCommit_Python_ValidatesContainer_Dataflow.yml/badge.svg?event=schedule)](https://github.com/apache/beam/actions/workflows/beam_PostCommit_Python_ValidatesContainer_Dataflow.yml?query=event%3Aschedule) | | [ PostCommit Python ValidatesContainer Dataflow With RC ](https://github.com/apache/beam/actions/workflows/beam_PostCommit_Python_ValidatesContainer_Dataflow_With_RC.yml) | ['3.9','3.10','3.11','3.12'] |`beam_PostCommit_Python_ValidatesContainer_Dataflow_With_RC.json`| [![.github/workflows/beam_PostCommit_Python_ValidatesContainer_Dataflow_With_RC.yml](https://github.com/apache/beam/actions/workflows/beam_PostCommit_Python_ValidatesContainer_Dataflow_With_RC.yml/badge.svg?event=schedule)](https://github.com/apache/beam/actions/workflows/beam_PostCommit_Python_ValidatesContainer_Dataflow_With_RC.yml?query=event%3Aschedule) | | [ PostCommit Python ValidatesRunner Dataflow ](https://github.com/apache/beam/actions/workflows/beam_PostCommit_Python_ValidatesRunner_Dataflow.yml) | ['3.9','3.11','3.12'] |`beam_PostCommit_Python_ValidatesRunner_Dataflow.json`| [![.github/workflows/beam_PostCommit_Python_ValidatesRunner_Dataflow.yml](https://github.com/apache/beam/actions/workflows/beam_PostCommit_Python_ValidatesRunner_Dataflow.yml/badge.svg?event=schedule)](https://github.com/apache/beam/actions/workflows/beam_PostCommit_Python_ValidatesRunner_Dataflow.yml?query=event%3Aschedule) | diff --git a/.github/workflows/beam_PostCommit_Java_Tpcds_Dataflow.yml b/.github/workflows/beam_PostCommit_Java_Tpcds_Dataflow.yml index 35bd18a7d541..e5bc892bb2ec 100644 --- a/.github/workflows/beam_PostCommit_Java_Tpcds_Dataflow.yml +++ b/.github/workflows/beam_PostCommit_Java_Tpcds_Dataflow.yml @@ -17,7 +17,7 @@ name: PostCommit Java Tpcds Dataflow on: schedule: - - cron: '30 4/6 * * *' + - cron: '30 4/8 * * *' pull_request: paths: ['release/trigger_all_tests.json', '.github/trigger_files/beam_PostCommit_Java_Tpcds_Dataflow.json'] workflow_dispatch: @@ -54,8 +54,8 @@ env: --dataSize=1GB --sourceType=PARQUET --dataDirectory=gs://beam-tpcds/datasets/parquet/nonpartitioned - --resultsDirectory=gs://beam-tpcds/results/dataflow/ - --tpcParallel=1 + --resultsDirectory=gs://temp-storage-for-perf-tests/tpcds/results/dataflow + --tpcParallel=4 INFLUXDB_USER: ${{ secrets.INFLUXDB_USER }} INFLUXDB_USER_PASSWORD: ${{ secrets.INFLUXDB_USER_PASSWORD }} tpcdsBigQueryArgs: | @@ -71,7 +71,7 @@ env: --baseInfluxMeasurement=tpcds --exportSummaryToInfluxDB=true --influxRetentionPolicy=forever - tpcdsQueriesArg: 3,7,10,25,26,29,35,38,40,42,43,52,55,69,79,83,84,87,93,96 + tpcdsQueriesArg: supported jobs: beam_PostCommit_Java_Tpcds_Dataflow: diff --git a/.github/workflows/beam_PostCommit_Java_Tpcds_Flink.yml b/.github/workflows/beam_PostCommit_Java_Tpcds_Flink.yml index f84dea31c345..bb4c3066ddf2 100644 --- a/.github/workflows/beam_PostCommit_Java_Tpcds_Flink.yml +++ b/.github/workflows/beam_PostCommit_Java_Tpcds_Flink.yml @@ -51,7 +51,7 @@ env: --dataSize=1GB --sourceType=PARQUET --dataDirectory=gs://beam-tpcds/datasets/parquet/nonpartitioned - --resultsDirectory=gs://beam-tpcds/results/flink/ + --resultsDirectory=gs://temp-storage-for-perf-tests/tpcds/results/flink --tpcParallel=1 INFLUXDB_USER: ${{ secrets.INFLUXDB_USER }} INFLUXDB_USER_PASSWORD: ${{ secrets.INFLUXDB_USER_PASSWORD }} @@ -68,7 +68,7 @@ env: --baseInfluxMeasurement=tpcds --exportSummaryToInfluxDB=true --influxRetentionPolicy=forever - tpcdsQueriesArg: 3,7,10,25,26,29,35,38,40,42,43,52,55,69,79,83,84,87,93,96 + tpcdsQueriesArg: supported jobs: beam_PostCommit_Java_Tpcds_Flink: diff --git a/.github/workflows/beam_PostCommit_Java_Tpcds_Spark.yml b/.github/workflows/beam_PostCommit_Java_Tpcds_Spark.yml index 6ce87271a4b9..8f67af866360 100644 --- a/.github/workflows/beam_PostCommit_Java_Tpcds_Spark.yml +++ b/.github/workflows/beam_PostCommit_Java_Tpcds_Spark.yml @@ -49,7 +49,7 @@ env: --dataSize=1GB --sourceType=PARQUET --dataDirectory=gs://beam-tpcds/datasets/parquet/nonpartitioned - --resultsDirectory=gs://beam-tpcds/results/spark3-rdd/ + --resultsDirectory=gs://temp-storage-for-perf-tests/tpcds/results/spark3-rdd --tpcParallel=1 INFLUXDB_USER: ${{ secrets.INFLUXDB_USER }} INFLUXDB_USER_PASSWORD: ${{ secrets.INFLUXDB_USER_PASSWORD }} @@ -66,7 +66,7 @@ env: --baseInfluxMeasurement=tpcds --exportSummaryToInfluxDB=true --influxRetentionPolicy=forever - tpcdsQueriesArg: 3,7,10,25,26,29,35,38,40,42,43,52,55,69,79,83,84,87,93,96 + tpcdsQueriesArg: supported jobs: beam_PostCommit_Java_Tpcds_Spark: diff --git a/.github/workflows/beam_PostCommit_Python_Tpcds_Prism.yml b/.github/workflows/beam_PostCommit_Python_Tpcds_Prism.yml new file mode 100644 index 000000000000..168d60b65f3d --- /dev/null +++ b/.github/workflows/beam_PostCommit_Python_Tpcds_Prism.yml @@ -0,0 +1,75 @@ +# Licensed to the Apache Software Foundation (ASF) under one or more +# contributor license agreements. See the NOTICE file distributed with +# this work for additional information regarding copyright ownership. +# The ASF licenses this file to You under the Apache License, Version 2.0 +# (the "License"); you may not use this file except in compliance with +# the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +name: PostCommit Python Tpcds Prism + +on: + schedule: + - cron: '30 4/6 * * *' + pull_request: + paths: ['release/trigger_all_tests.json', '.github/trigger_files/beam_PostCommit_Python_Tpcds_Prism.json'] + workflow_dispatch: + +# Setting explicit permissions for the action +permissions: + actions: write + pull-requests: read + checks: read + contents: read + deployments: read + id-token: none + issues: read + discussions: read + packages: read + pages: read + repository-projects: read + security-events: read + statuses: read + +# This allows a subsequently queued workflow run to interrupt previous runs +concurrency: + group: '${{ github.workflow }} @ ${{ github.event.pull_request.number || github.sha || github.head_ref || github.ref }}-${{ github.event.schedule || github.event.sender.login }}' + cancel-in-progress: true + +env: + DEVELOCITY_ACCESS_KEY: ${{ secrets.BEAM_DEVELOCITY_ACCESS_KEY }} + +jobs: + beam_PostCommit_Python_Tpcds_Prism: + if: | + github.event_name == 'workflow_dispatch' || + github.event_name == 'pull_request' || + (github.event_name == 'schedule' && github.repository == 'apache/beam') + runs-on: [self-hosted, ubuntu-24.04, main] + timeout-minutes: 60 + name: ${{ matrix.job_name }} + strategy: + matrix: + job_name: ["beam_PostCommit_Python_Tpcds_Prism"] + steps: + - uses: actions/checkout@v7 + with: + persist-credentials: false + - name: Setup repository + uses: ./.github/actions/setup-action + - name: Setup environment + uses: ./.github/actions/setup-environment-action + with: + python-version: default + go-version: default + - name: Run full TPC-DS suite on PrismRunner with regression check + uses: ./.github/actions/gradle-command-self-hosted-action + with: + gradle-command: :sdks:python:apache_beam:testing:benchmarks:tpcds:checkAllQueries diff --git a/sdks/java/extensions/sql/src/main/java/org/apache/beam/sdk/extensions/sql/impl/CalciteQueryPlanner.java b/sdks/java/extensions/sql/src/main/java/org/apache/beam/sdk/extensions/sql/impl/CalciteQueryPlanner.java index cedfc5dc6d1d..729ab8f8d84c 100644 --- a/sdks/java/extensions/sql/src/main/java/org/apache/beam/sdk/extensions/sql/impl/CalciteQueryPlanner.java +++ b/sdks/java/extensions/sql/src/main/java/org/apache/beam/sdk/extensions/sql/impl/CalciteQueryPlanner.java @@ -239,6 +239,11 @@ public BeamRelNode convertToBeamRel(String sqlStatement, QueryParameters queryPa String.format("Unable to convert query %s", sqlStatement), e); } catch (SqlParseException | ValidationException e) { throw new ParseException(String.format("Unable to parse query %s", sqlStatement), e); + } catch (RuntimeException e) { + if (e.getCause() instanceof SqlConversionException) { + throw new SqlConversionException(e.getMessage(), e.getCause()); + } + throw e; } finally { planner.close(); } diff --git a/sdks/java/extensions/sql/src/main/java/org/apache/beam/sdk/extensions/sql/impl/planner/BeamRuleSets.java b/sdks/java/extensions/sql/src/main/java/org/apache/beam/sdk/extensions/sql/impl/planner/BeamRuleSets.java index 8382fcb6e382..2e3ef1823863 100644 --- a/sdks/java/extensions/sql/src/main/java/org/apache/beam/sdk/extensions/sql/impl/planner/BeamRuleSets.java +++ b/sdks/java/extensions/sql/src/main/java/org/apache/beam/sdk/extensions/sql/impl/planner/BeamRuleSets.java @@ -26,6 +26,7 @@ import org.apache.beam.sdk.extensions.sql.impl.rule.BeamCalcMergeRule; import org.apache.beam.sdk.extensions.sql.impl.rule.BeamCalcRule; import org.apache.beam.sdk.extensions.sql.impl.rule.BeamCoGBKJoinRule; +import org.apache.beam.sdk.extensions.sql.impl.rule.BeamCrossJoinRule; import org.apache.beam.sdk.extensions.sql.impl.rule.BeamEnumerableConverterRule; import org.apache.beam.sdk.extensions.sql.impl.rule.BeamIOPushDownRule; import org.apache.beam.sdk.extensions.sql.impl.rule.BeamIntersectRule; @@ -60,6 +61,7 @@ public class BeamRuleSets { ImmutableList.of( // Rules for window functions CoreRules.PROJECT_TO_LOGICAL_PROJECT_AND_WINDOW, + CoreRules.CALC_TO_WINDOW, // Rules so we only have to implement Calc CoreRules.FILTER_TO_CALC, CoreRules.PROJECT_TO_CALC, @@ -144,6 +146,7 @@ public class BeamRuleSets { BeamUnnestRule.INSTANCE, BeamSideInputJoinRule.INSTANCE, BeamCoGBKJoinRule.INSTANCE, + BeamCrossJoinRule.INSTANCE, BeamSideInputLookupJoinRule.INSTANCE, BeamMatchRule.INSTANCE); diff --git a/sdks/java/extensions/sql/src/main/java/org/apache/beam/sdk/extensions/sql/impl/planner/RelMdNodeStats.java b/sdks/java/extensions/sql/src/main/java/org/apache/beam/sdk/extensions/sql/impl/planner/RelMdNodeStats.java index 40db8b074efa..51b197a7eac4 100644 --- a/sdks/java/extensions/sql/src/main/java/org/apache/beam/sdk/extensions/sql/impl/planner/RelMdNodeStats.java +++ b/sdks/java/extensions/sql/src/main/java/org/apache/beam/sdk/extensions/sql/impl/planner/RelMdNodeStats.java @@ -24,6 +24,7 @@ import org.apache.beam.sdk.extensions.sql.impl.rel.BeamRelNode; import org.apache.beam.vendor.calcite.v1_40_0.com.google.common.collect.Table; import org.apache.beam.vendor.calcite.v1_40_0.org.apache.calcite.rel.RelNode; +import org.apache.beam.vendor.calcite.v1_40_0.org.apache.calcite.rel.core.Values; import org.apache.beam.vendor.calcite.v1_40_0.org.apache.calcite.rel.metadata.MetadataDef; import org.apache.beam.vendor.calcite.v1_40_0.org.apache.calcite.rel.metadata.MetadataHandler; import org.apache.beam.vendor.calcite.v1_40_0.org.apache.calcite.rel.metadata.ReflectiveRelMetadataProvider; @@ -57,6 +58,11 @@ public NodeStats getNodeStats(RelNode rel, RelMetadataQuery mq) { return this.getBeamNodeStats((BeamRelNode) rel, bmq); } + if (rel instanceof Values) { + Values values = (Values) rel; + return NodeStats.create(values.getTuples().size()); + } + // We can later define custom methods for all different RelNodes to prevent hitting this point. // Similar to RelMdRowCount in calcite. diff --git a/sdks/java/extensions/sql/src/main/java/org/apache/beam/sdk/extensions/sql/impl/rel/BeamCrossJoinRel.java b/sdks/java/extensions/sql/src/main/java/org/apache/beam/sdk/extensions/sql/impl/rel/BeamCrossJoinRel.java new file mode 100644 index 000000000000..ebdc73796887 --- /dev/null +++ b/sdks/java/extensions/sql/src/main/java/org/apache/beam/sdk/extensions/sql/impl/rel/BeamCrossJoinRel.java @@ -0,0 +1,173 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package org.apache.beam.sdk.extensions.sql.impl.rel; + +import java.util.Collections; +import java.util.List; +import java.util.Set; +import org.apache.beam.sdk.extensions.sql.impl.utils.CalciteUtils; +import org.apache.beam.sdk.schemas.Schema; +import org.apache.beam.sdk.transforms.DoFn; +import org.apache.beam.sdk.transforms.PTransform; +import org.apache.beam.sdk.transforms.ParDo; +import org.apache.beam.sdk.transforms.View; +import org.apache.beam.sdk.transforms.windowing.TimestampCombiner; +import org.apache.beam.sdk.transforms.windowing.Window; +import org.apache.beam.sdk.values.PCollection; +import org.apache.beam.sdk.values.PCollectionList; +import org.apache.beam.sdk.values.PCollectionView; +import org.apache.beam.sdk.values.Row; +import org.apache.beam.vendor.calcite.v1_40_0.org.apache.calcite.plan.RelOptCluster; +import org.apache.beam.vendor.calcite.v1_40_0.org.apache.calcite.plan.RelTraitSet; +import org.apache.beam.vendor.calcite.v1_40_0.org.apache.calcite.rel.RelNode; +import org.apache.beam.vendor.calcite.v1_40_0.org.apache.calcite.rel.core.CorrelationId; +import org.apache.beam.vendor.calcite.v1_40_0.org.apache.calcite.rel.core.Join; +import org.apache.beam.vendor.calcite.v1_40_0.org.apache.calcite.rel.core.JoinRelType; +import org.apache.beam.vendor.calcite.v1_40_0.org.apache.calcite.rex.RexNode; + +/** + * A {@code BeamJoinRel} that implements a BOUNDED cartesian (CROSS) join, used for the literal-TRUE + * join conditions that Calcite emits when it decorrelates uncorrelated scalar / IN / EXISTS + * subqueries and lateral joins. + * + *

Only supported when BOTH inputs are bounded. The right side is broadcast as a side input and + * the left side is iterated; every (left, right) pair is emitted as one combined row (left values + * followed by right values). + * + *

First increment: INNER join with an always-true condition. A non-trivial (residual) condition + * is not yet evaluated here. + */ +public class BeamCrossJoinRel extends BeamJoinRel { + + public BeamCrossJoinRel( + RelOptCluster cluster, + RelTraitSet traitSet, + RelNode left, + RelNode right, + RexNode condition, + Set variablesSet, + JoinRelType joinType) { + super(cluster, traitSet, left, right, condition, variablesSet, joinType); + } + + @Override + public Join copy( + RelTraitSet traitSet, + RexNode conditionExpr, + RelNode left, + RelNode right, + JoinRelType joinType, + boolean semiJoinDone) { + return new BeamCrossJoinRel( + getCluster(), traitSet, left, right, conditionExpr, variablesSet, joinType); + } + + @Override + public PTransform, PCollection> buildPTransform() { + BeamRelNode leftRelNode = BeamSqlRelUtils.getBeamRelInput(left); + BeamRelNode rightRelNode = BeamSqlRelUtils.getBeamRelInput(right); + if (leftRelNode.isBounded() != PCollection.IsBounded.BOUNDED + || rightRelNode.isBounded() != PCollection.IsBounded.BOUNDED) { + throw new UnsupportedOperationException( + "Bounded cross join requires both inputs to be bounded."); + } + if (joinType != JoinRelType.INNER && joinType != JoinRelType.LEFT) { + throw new UnsupportedOperationException( + "Only INNER and LEFT cross join are supported for now, got " + joinType.name()); + } + if (!condition.isAlwaysTrue()) { + throw new UnsupportedOperationException( + "Cross join with a residual (non-equi) condition is not supported yet."); + } + return new CrossJoin(); + } + + private class CrossJoin extends PTransform, PCollection> { + @Override + public PCollection expand(PCollectionList pinput) { + PCollection leftRows = + pinput + .get(0) + .apply( + "left_TimestampCombiner", + Window.configure().withTimestampCombiner(TimestampCombiner.EARLIEST)); + PCollection rightRows = + pinput + .get(1) + .apply( + "right_TimestampCombiner", + Window.configure().withTimestampCombiner(TimestampCombiner.EARLIEST)); + + final Schema outputSchema = CalciteUtils.toSchema(getRowType()); + final PCollectionView> rightView = rightRows.apply(View.asList()); + final boolean leftJoin = joinType == JoinRelType.LEFT; + final int rightFieldCount = getRowType().getFieldCount() - left.getRowType().getFieldCount(); + + return leftRows + .apply( + ParDo.of(new CrossJoinDoFn(outputSchema, rightView, leftJoin, rightFieldCount)) + .withSideInputs(rightView)) + .setRowSchema(outputSchema); + } + } + + /** + * Emits the cartesian product of each left row with the broadcast right rows. A static {@link + * DoFn} (not an anonymous/inner class) so it does not capture the non-serializable enclosing rel. + */ + private static class CrossJoinDoFn extends DoFn { + private final Schema outputSchema; + private final PCollectionView> rightView; + private final boolean leftJoin; + private final int rightFieldCount; + + CrossJoinDoFn( + Schema outputSchema, + PCollectionView> rightView, + boolean leftJoin, + int rightFieldCount) { + this.outputSchema = outputSchema; + this.rightView = rightView; + this.leftJoin = leftJoin; + this.rightFieldCount = rightFieldCount; + } + + // Row.Builder.addValues expects List<@Nullable Object>, while Row.getBaseValues() returns + // List<@NonNull Object>, so it is safe to suppress nullness here. + @SuppressWarnings("nullness") + @ProcessElement + public void processElement(ProcessContext c) { + Row leftRow = c.element(); + List rightRowList = c.sideInput(rightView); + for (Row rightRow : rightRowList) { + c.output( + Row.withSchema(outputSchema) + .addValues(leftRow.getBaseValues()) + .addValues(rightRow.getBaseValues()) + .build()); + } + if (leftJoin && rightRowList.isEmpty()) { + c.output( + Row.withSchema(outputSchema) + .addValues(leftRow.getBaseValues()) + .addValues(Collections.nCopies(rightFieldCount, null)) + .build()); + } + } + } +} diff --git a/sdks/java/extensions/sql/src/main/java/org/apache/beam/sdk/extensions/sql/impl/rel/BeamSetOperatorRelBase.java b/sdks/java/extensions/sql/src/main/java/org/apache/beam/sdk/extensions/sql/impl/rel/BeamSetOperatorRelBase.java index cad845c2b3a9..efbdfc4111b5 100644 --- a/sdks/java/extensions/sql/src/main/java/org/apache/beam/sdk/extensions/sql/impl/rel/BeamSetOperatorRelBase.java +++ b/sdks/java/extensions/sql/src/main/java/org/apache/beam/sdk/extensions/sql/impl/rel/BeamSetOperatorRelBase.java @@ -17,12 +17,19 @@ */ package org.apache.beam.sdk.extensions.sql.impl.rel; -import static org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.Preconditions.checkArgument; - import java.io.Serializable; +import java.util.ArrayList; +import java.util.Collections; +import java.util.List; import org.apache.beam.sdk.extensions.sql.impl.transform.BeamSetOperatorsTransforms; +import org.apache.beam.sdk.extensions.sql.impl.utils.CalciteUtils; +import org.apache.beam.sdk.schemas.Schema; +import org.apache.beam.sdk.schemas.Schema.FieldType; +import org.apache.beam.sdk.schemas.Schema.TypeName; import org.apache.beam.sdk.schemas.transforms.CoGroup; import org.apache.beam.sdk.schemas.transforms.CoGroup.By; +import org.apache.beam.sdk.transforms.Flatten; +import org.apache.beam.sdk.transforms.MapElements; import org.apache.beam.sdk.transforms.PTransform; import org.apache.beam.sdk.transforms.ParDo; import org.apache.beam.sdk.transforms.windowing.WindowFn; @@ -30,6 +37,7 @@ import org.apache.beam.sdk.values.PCollectionList; import org.apache.beam.sdk.values.PCollectionTuple; import org.apache.beam.sdk.values.Row; +import org.apache.beam.sdk.values.TypeDescriptors; /** * Delegate for Set operators: {@code BeamUnionRel}, {@code BeamIntersectRel} and {@code @@ -59,13 +67,31 @@ public BeamSetOperatorRelBase(BeamRelNode beamRelNode, OpType opType, boolean al @Override public PCollection expand(PCollectionList inputs) { - checkArgument( - inputs.size() == 2, - "Wrong number of arguments to %s: %s", - beamRelNode.getClass().getSimpleName(), - inputs); - PCollection leftRows = inputs.get(0); - PCollection rightRows = inputs.get(1); + Schema targetSchema = CalciteUtils.toSchema(beamRelNode.getRowType()); + + if (opType == OpType.UNION && all) { + PCollectionList coercedInputs = PCollectionList.empty(inputs.getPipeline()); + for (PCollection input : inputs.getAll()) { + coercedInputs = coercedInputs.and(coerce(input, targetSchema)); + } + return coercedInputs.apply("UnionAllFlatten", Flatten.pCollections()); + } + + if (inputs.size() == 2) { + return expandPair(coerce(inputs.get(0), targetSchema), coerce(inputs.get(1), targetSchema)); + } else if (inputs.size() > 2) { + PCollection result = coerce(inputs.get(0), targetSchema); + for (int i = 1; i < inputs.size(); i++) { + result = expandPair(result, coerce(inputs.get(i), targetSchema)); + } + return result; + } else { + throw new IllegalArgumentException( + "Too few arguments to " + beamRelNode.getClass().getSimpleName()); + } + } + + private PCollection expandPair(PCollection leftRows, PCollection rightRows) { WindowFn leftWindow = leftRows.getWindowingStrategy().getWindowFn(); WindowFn rightWindow = rightRows.getWindowingStrategy().getWindowFn(); @@ -95,4 +121,54 @@ public PCollection expand(PCollectionList inputs) { lhsTag, rhsTag, opType, all))) .setRowSchema(joined.getSchema().getField("key").getType().getRowSchema()); } + + private static PCollection coerce(PCollection input, Schema targetSchema) { + if (input.getSchema().equals(targetSchema)) { + return input; + } + return input + .apply( + "CoerceSchema", + MapElements.into(TypeDescriptors.rows()).via(row -> coerceRow(row, targetSchema))) + .setRowSchema(targetSchema); + } + + private static Row coerceRow(Row row, Schema targetSchema) { + Row.Builder builder = Row.withSchema(targetSchema); + for (int i = 0; i < targetSchema.getFieldCount(); i++) { + Schema.Field targetField = targetSchema.getField(i); + Object value = row.getValue(i); + builder.addValue( + coerceValue(value, row.getSchema().getField(i).getType(), targetField.getType())); + } + return builder.build(); + } + + private static Object coerceValue(Object value, FieldType inputType, FieldType targetType) { + if (value == null) { + return null; + } + if (inputType.equals(targetType)) { + return value; + } + if (inputType.getTypeName() == TypeName.ARRAY && targetType.getTypeName() == TypeName.ARRAY) { + List list = (List) value; + if (list.isEmpty()) { + return Collections.emptyList(); + } + List coercedList = new ArrayList<>(list.size()); + for (Object element : list) { + coercedList.add( + coerceValue( + element, + inputType.getCollectionElementType(), + targetType.getCollectionElementType())); + } + return coercedList; + } + if (inputType.getTypeName() == TypeName.ROW && targetType.getTypeName() == TypeName.ROW) { + return coerceRow((Row) value, targetType.getRowSchema()); + } + return value; + } } diff --git a/sdks/java/extensions/sql/src/main/java/org/apache/beam/sdk/extensions/sql/impl/rel/BeamSortRel.java b/sdks/java/extensions/sql/src/main/java/org/apache/beam/sdk/extensions/sql/impl/rel/BeamSortRel.java index c7a0b9136bbb..a338d6b2e8b6 100644 --- a/sdks/java/extensions/sql/src/main/java/org/apache/beam/sdk/extensions/sql/impl/rel/BeamSortRel.java +++ b/sdks/java/extensions/sql/src/main/java/org/apache/beam/sdk/extensions/sql/impl/rel/BeamSortRel.java @@ -378,9 +378,12 @@ public int compare(Row row1, Row row2) { if (isValue1Null && isValue2Null) { continue; } else if (isValue1Null && !isValue2Null) { - fieldRet = -1 * (nullsFirst.get(i) ? -1 : 1); + // Return early as NULLS FIRST/LAST is independent of ASC/DESC orientation. + fieldRet = nullsFirst.get(i) ? -1 : 1; + return fieldRet; } else if (!isValue1Null && isValue2Null) { - fieldRet = 1 * (nullsFirst.get(i) ? -1 : 1); + fieldRet = nullsFirst.get(i) ? 1 : -1; + return fieldRet; } else { switch (sqlTypeName) { case TINYINT: @@ -389,8 +392,12 @@ public int compare(Row row1, Row row2) { case BIGINT: case FLOAT: case DOUBLE: + case DECIMAL: + case BOOLEAN: case VARCHAR: + case CHAR: case DATE: + case TIME: case TIMESTAMP: Comparable v1 = row1.getBaseValue(fieldIndex, Comparable.class); Comparable v2 = row2.getBaseValue(fieldIndex, Comparable.class); diff --git a/sdks/java/extensions/sql/src/main/java/org/apache/beam/sdk/extensions/sql/impl/rel/BeamWindowRel.java b/sdks/java/extensions/sql/src/main/java/org/apache/beam/sdk/extensions/sql/impl/rel/BeamWindowRel.java index e597b6c63ed4..6fa4e383a2be 100644 --- a/sdks/java/extensions/sql/src/main/java/org/apache/beam/sdk/extensions/sql/impl/rel/BeamWindowRel.java +++ b/sdks/java/extensions/sql/src/main/java/org/apache/beam/sdk/extensions/sql/impl/rel/BeamWindowRel.java @@ -54,6 +54,7 @@ import org.apache.beam.vendor.calcite.v1_40_0.org.apache.calcite.rex.RexLiteral; import org.apache.beam.vendor.calcite.v1_40_0.org.apache.calcite.rex.RexNode; import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.Lists; +import org.joda.time.ReadableInstant; /** * {@code BeamRelNode} to replace a {@code Window} node. @@ -142,9 +143,27 @@ public PTransform, PCollection> buildPTransform() { List argList = anAggCall.getArgList(); Schema.Field field = CalciteUtils.toField(anAggCall.getName(), anAggCall.getType()); - Combine.CombineFn combineFn = - AggregationCombineFnAdapter.createCombineFnAnalyticsFunctions( - anAggCall, field, anAggCall.getAggregation().getName()); + String aggName = anAggCall.getAggregation().getName(); + Combine.CombineFn combineFn; + if (anAggCall.ignoreNulls() + && ("FIRST_VALUE".equals(aggName) || "LAST_VALUE".equals(aggName))) { + combineFn = + "FIRST_VALUE".equals(aggName) + ? BeamBuiltinAnalyticFunctions.navigationFirstValue(true) + : BeamBuiltinAnalyticFunctions.navigationLastValue(true); + } else if ("NTH_VALUE".equals(aggName)) { + long nthN = readNthValueOffset(argList); + combineFn = + BeamBuiltinAnalyticFunctions.navigationNthValue( + nthN, anAggCall.ignoreNulls()); + // Keep only the value column in argList; the offset n is resolved from + // constants. + argList = argList.isEmpty() ? argList : argList.subList(0, 1); + } else { + combineFn = + AggregationCombineFnAdapter.createCombineFnAnalyticsFunctions( + anAggCall, field, aggName); + } FieldAggregation fieldAggregation = new FieldAggregation( partitionKeysDef, @@ -169,6 +188,31 @@ private BigDecimal getLiteralValueConstants(RexNode n) { return (BigDecimal) this.constants.get(idx).getValue(); } + /** + * Extracts the 1-based position {@code n} of an {@code NTH_VALUE(expr, n)} call. The value column + * is {@code argList[0]}; {@code argList[1]} references the {@code n} literal, which the planner + * stores in the Window {@code constants} list at index {@code arg - inputFieldCount}. Defaults to + * 1 (== FIRST_VALUE) if the position cannot be resolved. + */ + private long readNthValueOffset(List argList) { + if (argList.size() < 2) { + return 1L; + } + int inputFieldCount = input.getRowType().getFieldCount(); + // Calcite stores window literal arguments in `constants` indexed past the input fields. + int constIdx = argList.get(1) - inputFieldCount; + if (constIdx >= 0 && constIdx < this.constants.size()) { + Object v = this.constants.get(constIdx).getValue(); + if (v instanceof BigDecimal) { + return ((BigDecimal) v).longValueExact(); + } + if (v instanceof Number) { + return ((Number) v).longValue(); + } + } + return 1L; + } + private static class FieldAggregation implements Serializable { private List partitionKeys; @@ -208,8 +252,9 @@ public FieldAggregation( @Override public NodeStats estimateNodeStats(BeamRelMetadataQuery mq) { - NodeStats inputStat = BeamSqlRelUtils.getNodeStats(this.input, mq); - return inputStat; + // Use mq.getRowCount so RelSubset inputs do not return UNKNOWN/infinite cost. + double rows = mq.getRowCount(this.input); + return NodeStats.create(rows, 0d, rows); } /** @@ -220,10 +265,12 @@ public NodeStats estimateNodeStats(BeamRelMetadataQuery mq) { */ @Override public BeamCostModel beamComputeSelfCost(RelOptPlanner planner, BeamRelMetadataQuery mq) { - NodeStats inputStat = BeamSqlRelUtils.getNodeStats(this.input, mq); + // Use the planner's row-count (RelSubset-safe) rather than BeamSqlRelUtils.getNodeStats, which + // yields UNKNOWN/infinite for a RelSubset input and prevents the Volcano cost fixpoint from + // converging. + double rows = mq.getRowCount(this.input); float multiplier = 1f + 0.125f; - return BeamCostModel.FACTORY.makeCost( - inputStat.getRowCount() * multiplier, inputStat.getRate() * multiplier); + return BeamCostModel.FACTORY.makeCost(rows * multiplier, 0d); } private static class Transform extends PTransform, PCollection> { @@ -309,16 +356,21 @@ public void processElement( if (fieldAgg.combineFn instanceof BeamBuiltinAnalyticFunctions.PositionAwareCombineFn) { BeamBuiltinAnalyticFunctions.PositionAwareCombineFn fn = (BeamBuiltinAnalyticFunctions.PositionAwareCombineFn) fieldAgg.combineFn; + Object positionInputValue = + aggFieldIndex >= 0 ? aggRow.getBaseValue(aggFieldIndex) : (Object) aggRow; accumulator = fn.addInput( accumulator, getOrderByValue(aggRow), + positionInputValue, count, (long) idx, (long) sortedRowsAsList.size()); } else { - accumulator = - fieldAgg.combineFn.addInput(accumulator, aggRow.getBaseValue(aggFieldIndex)); + // Pass the row itself for argumentless functions like COUNT(*) OVER (...). + Object inputValue = + aggFieldIndex >= 0 ? aggRow.getBaseValue(aggFieldIndex) : (Object) aggRow; + accumulator = fieldAgg.combineFn.addInput(accumulator, inputValue); } count++; } @@ -384,8 +436,33 @@ private BigDecimal getOrderByValue(Row r) { if (fieldAgg.orderKeys.size() == 0) { return null; } else { - return new BigDecimal(((Number) r.getBaseValue(fieldAgg.orderKeys.get(0))).toString()); + return toComparableBigDecimal(r.getBaseValue(fieldAgg.orderKeys.get(0))); + } + } + + /** + * Convert an ORDER BY key's base value into a comparable {@link BigDecimal} for RANGE + * framing. + * + *

Numeric values keep their exact decimal representation (byte-identical to the historical + * {@code new BigDecimal(((Number) v).toString())} path). Temporal values (DATE / TIME / + * TIMESTAMP / TIMESTAMP_WITH_LOCAL_TIME_ZONE) all surface a base value of joda {@link + * ReadableInstant}; their epoch-millis are used so a RANGE frame ordered by a temporal column + * compares correctly instead of throwing a {@link ClassCastException}. A null value (NULL + * ordering key) maps to {@code null}, leaving the surrounding callers' existing null handling + * intact. + */ + private BigDecimal toComparableBigDecimal(Object value) { + if (value == null) { + return null; + } + if (value instanceof Number) { + return new BigDecimal(((Number) value).toString()); + } + if (value instanceof ReadableInstant) { + return BigDecimal.valueOf(((ReadableInstant) value).getMillis()); } + return BigDecimal.ZERO; } private List getRows(List input, int index) { @@ -395,8 +472,9 @@ private List getRows(List input, int index) { fieldAgg.upperLimit != null ? fieldAgg.upperLimit.intValue() : Integer.MAX_VALUE; int lowerIndex = ll == Integer.MAX_VALUE ? Integer.MIN_VALUE : index - ll; int upperIndex = ul == Integer.MAX_VALUE ? Integer.MAX_VALUE : index + ul + 1; - lowerIndex = lowerIndex < 0 ? 0 : lowerIndex; - upperIndex = upperIndex > input.size() ? input.size() : upperIndex; + // Clamp bounds to [0, size] so frames completely outside the partition yield an empty list. + lowerIndex = Math.max(0, Math.min(lowerIndex, input.size())); + upperIndex = Math.max(lowerIndex, Math.min(upperIndex, input.size())); List out = input.subList(lowerIndex, upperIndex); return out; } diff --git a/sdks/java/extensions/sql/src/main/java/org/apache/beam/sdk/extensions/sql/impl/rule/BeamBasicAggregationRule.java b/sdks/java/extensions/sql/src/main/java/org/apache/beam/sdk/extensions/sql/impl/rule/BeamBasicAggregationRule.java index e2df3c3a685f..002477a717b8 100644 --- a/sdks/java/extensions/sql/src/main/java/org/apache/beam/sdk/extensions/sql/impl/rule/BeamBasicAggregationRule.java +++ b/sdks/java/extensions/sql/src/main/java/org/apache/beam/sdk/extensions/sql/impl/rule/BeamBasicAggregationRule.java @@ -20,6 +20,7 @@ import java.util.ArrayList; import java.util.List; import java.util.stream.Collectors; +import org.apache.beam.sdk.extensions.sql.impl.SqlConversionException; import org.apache.beam.sdk.extensions.sql.impl.rel.BeamAggregationRel; import org.apache.beam.sdk.extensions.sql.impl.rel.BeamLogicalConvention; import org.apache.beam.vendor.calcite.v1_40_0.org.apache.calcite.plan.RelOptRule; @@ -61,7 +62,8 @@ public void onMatch(RelOptRuleCall call) { RelNode relNode = call.rel(1); if (aggregate.getGroupType() != Aggregate.Group.SIMPLE) { - return; + throw new SqlConversionException( + "GROUP BY " + aggregate.getGroupType() + " is not supported."); } if (relNode instanceof Project || relNode instanceof Calc || relNode instanceof Filter) { diff --git a/sdks/java/extensions/sql/src/main/java/org/apache/beam/sdk/extensions/sql/impl/rule/BeamCoGBKJoinRule.java b/sdks/java/extensions/sql/src/main/java/org/apache/beam/sdk/extensions/sql/impl/rule/BeamCoGBKJoinRule.java index f44839de062d..156b3740e7b4 100644 --- a/sdks/java/extensions/sql/src/main/java/org/apache/beam/sdk/extensions/sql/impl/rule/BeamCoGBKJoinRule.java +++ b/sdks/java/extensions/sql/src/main/java/org/apache/beam/sdk/extensions/sql/impl/rule/BeamCoGBKJoinRule.java @@ -17,16 +17,28 @@ */ package org.apache.beam.sdk.extensions.sql.impl.rule; +import java.util.ArrayList; +import java.util.List; +import org.apache.beam.sdk.extensions.sql.impl.rel.BeamCalcRel; import org.apache.beam.sdk.extensions.sql.impl.rel.BeamCoGBKJoinRel; import org.apache.beam.sdk.extensions.sql.impl.rel.BeamJoinRel; import org.apache.beam.sdk.extensions.sql.impl.rel.BeamLogicalConvention; import org.apache.beam.sdk.values.PCollection; import org.apache.beam.vendor.calcite.v1_40_0.org.apache.calcite.plan.RelOptRule; import org.apache.beam.vendor.calcite.v1_40_0.org.apache.calcite.plan.RelOptRuleCall; +import org.apache.beam.vendor.calcite.v1_40_0.org.apache.calcite.plan.RelOptUtil; import org.apache.beam.vendor.calcite.v1_40_0.org.apache.calcite.rel.RelNode; import org.apache.beam.vendor.calcite.v1_40_0.org.apache.calcite.rel.core.Join; +import org.apache.beam.vendor.calcite.v1_40_0.org.apache.calcite.rel.core.JoinRelType; import org.apache.beam.vendor.calcite.v1_40_0.org.apache.calcite.rel.core.RelFactories; import org.apache.beam.vendor.calcite.v1_40_0.org.apache.calcite.rel.logical.LogicalJoin; +import org.apache.beam.vendor.calcite.v1_40_0.org.apache.calcite.rex.RexBuilder; +import org.apache.beam.vendor.calcite.v1_40_0.org.apache.calcite.rex.RexCall; +import org.apache.beam.vendor.calcite.v1_40_0.org.apache.calcite.rex.RexFieldAccess; +import org.apache.beam.vendor.calcite.v1_40_0.org.apache.calcite.rex.RexInputRef; +import org.apache.beam.vendor.calcite.v1_40_0.org.apache.calcite.rex.RexNode; +import org.apache.beam.vendor.calcite.v1_40_0.org.apache.calcite.rex.RexProgramBuilder; +import org.apache.beam.vendor.calcite.v1_40_0.org.apache.calcite.sql.fun.SqlStdOperatorTable; /** * Rule to convert {@code LogicalJoin} node to {@code BeamCoGBKJoinRel} node. @@ -55,6 +67,12 @@ public boolean matches(RelOptRuleCall call) { if (BeamJoinRel.containsSeekableInput(call.rel(0))) { return false; } + // Bounded pure cross joins are handled by BeamCrossJoinRule. + if (((Join) call.rel(0)).getCondition().isAlwaysTrue() + && BeamJoinRel.getBoundednessOfRelNode(call.rel(1)) == PCollection.IsBounded.BOUNDED + && BeamJoinRel.getBoundednessOfRelNode(call.rel(2)) == PCollection.IsBounded.BOUNDED) { + return false; + } PCollection.IsBounded boundednessOfLeftRel = BeamJoinRel.getBoundednessOfRelNode(call.rel(1)); PCollection.IsBounded boundednessOfRightRel = BeamJoinRel.getBoundednessOfRelNode(call.rel(2)); return (boundednessOfLeftRel == boundednessOfRightRel); @@ -64,6 +82,49 @@ public boolean matches(RelOptRuleCall call) { public void onMatch(RelOptRuleCall call) { Join join = (Join) call.rel(0); + RexNode condition = join.getCondition(); + List conjuncts = new ArrayList<>(); + RelOptUtil.decomposeConjunction(condition, conjuncts); + + List equiConditions = new ArrayList<>(); + List residualConditions = new ArrayList<>(); + int leftFieldCount = join.getLeft().getRowType().getFieldCount(); + + for (RexNode conjunct : conjuncts) { + if (isEquiJoinConjunct(conjunct, leftFieldCount)) { + equiConditions.add(conjunct); + } else { + residualConditions.add(conjunct); + } + } + + if (equiConditions.isEmpty() + || (!residualConditions.isEmpty() && join.getJoinType() != JoinRelType.INNER)) { + BeamCoGBKJoinRel rel = + new BeamCoGBKJoinRel( + join.getCluster(), + join.getTraitSet().replace(BeamLogicalConvention.INSTANCE), + convert( + join.getLeft(), + join.getLeft().getTraitSet().replace(BeamLogicalConvention.INSTANCE)), + convert( + join.getRight(), + join.getRight().getTraitSet().replace(BeamLogicalConvention.INSTANCE)), + condition, + join.getVariablesSet(), + join.getJoinType()); + call.transformTo(rel); + return; + } + + RexBuilder rexBuilder = join.getCluster().getRexBuilder(); + RexNode finalEquiCondition; + if (equiConditions.size() == 1) { + finalEquiCondition = equiConditions.get(0); + } else { + finalEquiCondition = rexBuilder.makeCall(SqlStdOperatorTable.AND, equiConditions); + } + BeamCoGBKJoinRel rel = new BeamCoGBKJoinRel( join.getCluster(), @@ -74,9 +135,63 @@ public void onMatch(RelOptRuleCall call) { convert( join.getRight(), join.getRight().getTraitSet().replace(BeamLogicalConvention.INSTANCE)), - join.getCondition(), + finalEquiCondition, join.getVariablesSet(), join.getJoinType()); - call.transformTo(rel); + + if (residualConditions.isEmpty()) { + call.transformTo(rel); + } else { + RexNode finalResidualCondition; + if (residualConditions.size() == 1) { + finalResidualCondition = residualConditions.get(0); + } else { + finalResidualCondition = rexBuilder.makeCall(SqlStdOperatorTable.AND, residualConditions); + } + RexProgramBuilder programBuilder = new RexProgramBuilder(rel.getRowType(), rexBuilder); + programBuilder.addIdentity(); + programBuilder.addCondition(finalResidualCondition); + BeamCalcRel calcRel = + new BeamCalcRel( + join.getCluster(), + join.getTraitSet().replace(BeamLogicalConvention.INSTANCE), + rel, + programBuilder.getProgram()); + call.transformTo(calcRel); + } + } + + private static int extractInputIndex(RexNode node) { + if (node instanceof RexInputRef) { + return ((RexInputRef) node).getIndex(); + } + if (node instanceof RexFieldAccess) { + return extractInputIndex(((RexFieldAccess) node).getReferenceExpr()); + } + return -1; + } + + private boolean isEquiJoinConjunct(RexNode conjunct, int leftFieldCount) { + if (!(conjunct instanceof RexCall)) { + return false; + } + RexCall call = (RexCall) conjunct; + if (!call.getOperator().getName().equals("=")) { + return false; + } + if (call.getOperands().size() != 2) { + return false; + } + RexNode op0 = call.getOperands().get(0); + RexNode op1 = call.getOperands().get(1); + + int idx0 = extractInputIndex(op0); + int idx1 = extractInputIndex(op1); + if (idx0 < 0 || idx1 < 0) { + return false; + } + boolean op0Left = idx0 < leftFieldCount; + boolean op1Left = idx1 < leftFieldCount; + return op0Left != op1Left; } } diff --git a/sdks/java/extensions/sql/src/main/java/org/apache/beam/sdk/extensions/sql/impl/rule/BeamCrossJoinRule.java b/sdks/java/extensions/sql/src/main/java/org/apache/beam/sdk/extensions/sql/impl/rule/BeamCrossJoinRule.java new file mode 100644 index 000000000000..18ceb9f0dd33 --- /dev/null +++ b/sdks/java/extensions/sql/src/main/java/org/apache/beam/sdk/extensions/sql/impl/rule/BeamCrossJoinRule.java @@ -0,0 +1,79 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package org.apache.beam.sdk.extensions.sql.impl.rule; + +import org.apache.beam.sdk.extensions.sql.impl.rel.BeamCrossJoinRel; +import org.apache.beam.sdk.extensions.sql.impl.rel.BeamJoinRel; +import org.apache.beam.sdk.extensions.sql.impl.rel.BeamLogicalConvention; +import org.apache.beam.sdk.values.PCollection; +import org.apache.beam.vendor.calcite.v1_40_0.org.apache.calcite.plan.RelOptRule; +import org.apache.beam.vendor.calcite.v1_40_0.org.apache.calcite.plan.RelOptRuleCall; +import org.apache.beam.vendor.calcite.v1_40_0.org.apache.calcite.rel.RelNode; +import org.apache.beam.vendor.calcite.v1_40_0.org.apache.calcite.rel.core.Join; +import org.apache.beam.vendor.calcite.v1_40_0.org.apache.calcite.rel.core.RelFactories; +import org.apache.beam.vendor.calcite.v1_40_0.org.apache.calcite.rel.logical.LogicalJoin; + +/** + * Rule to convert a {@code LogicalJoin} with a cross / non-equi condition into a {@code + * BeamCrossJoinRel}, but only when BOTH inputs are bounded. Equi-joins keep using the existing + * CoGBK / side-input rules; those rules are gated with {@code isJoinLegal} so they no longer match + * the cross / non-equi case. Unbounded cross joins are intentionally not matched (rejected). + */ +public class BeamCrossJoinRule extends RelOptRule { + public static final BeamCrossJoinRule INSTANCE = new BeamCrossJoinRule(); + + private BeamCrossJoinRule() { + super( + operand(LogicalJoin.class, operand(RelNode.class, any()), operand(RelNode.class, any())), + RelFactories.LOGICAL_BUILDER, + "BeamCrossJoinRule"); + } + + @Override + public boolean matches(RelOptRuleCall call) { + Join join = (Join) call.rel(0); + // Only match pure cross joins; non-equi joins stay on CoGBK/side-input rules. + if (!join.getCondition().isAlwaysTrue()) { + return false; + } + if (BeamJoinRel.containsSeekableInput(call.rel(0))) { + return false; + } + return BeamJoinRel.getBoundednessOfRelNode(call.rel(1)) == PCollection.IsBounded.BOUNDED + && BeamJoinRel.getBoundednessOfRelNode(call.rel(2)) == PCollection.IsBounded.BOUNDED; + } + + @Override + public void onMatch(RelOptRuleCall call) { + Join join = (Join) call.rel(0); + BeamCrossJoinRel rel = + new BeamCrossJoinRel( + join.getCluster(), + join.getTraitSet().replace(BeamLogicalConvention.INSTANCE), + convert( + join.getLeft(), + join.getLeft().getTraitSet().replace(BeamLogicalConvention.INSTANCE)), + convert( + join.getRight(), + join.getRight().getTraitSet().replace(BeamLogicalConvention.INSTANCE)), + join.getCondition(), + join.getVariablesSet(), + join.getJoinType()); + call.transformTo(rel); + } +} diff --git a/sdks/java/extensions/sql/src/main/java/org/apache/beam/sdk/extensions/sql/impl/rule/BeamWindowRule.java b/sdks/java/extensions/sql/src/main/java/org/apache/beam/sdk/extensions/sql/impl/rule/BeamWindowRule.java index c442ae3d1947..ad1d94bbc52f 100644 --- a/sdks/java/extensions/sql/src/main/java/org/apache/beam/sdk/extensions/sql/impl/rule/BeamWindowRule.java +++ b/sdks/java/extensions/sql/src/main/java/org/apache/beam/sdk/extensions/sql/impl/rule/BeamWindowRule.java @@ -36,10 +36,12 @@ private BeamWindowRule() { @Override public RelNode convert(RelNode relNode) { Window w = (Window) relNode; + RelNode input = w.getInput(); return new BeamWindowRel( w.getCluster(), w.getTraitSet().replace(BeamLogicalConvention.INSTANCE), - w.getInput(), + // Convert input to BeamLogicalConvention so inputs from CALC_TO_WINDOW can be planned. + convert(input, input.getTraitSet().replace(BeamLogicalConvention.INSTANCE)), w.getConstants(), w.getRowType(), w.groups); diff --git a/sdks/java/extensions/sql/src/main/java/org/apache/beam/sdk/extensions/sql/impl/schema/BeamTableUtils.java b/sdks/java/extensions/sql/src/main/java/org/apache/beam/sdk/extensions/sql/impl/schema/BeamTableUtils.java index d9dbf881850e..8eb23aef61fe 100644 --- a/sdks/java/extensions/sql/src/main/java/org/apache/beam/sdk/extensions/sql/impl/schema/BeamTableUtils.java +++ b/sdks/java/extensions/sql/src/main/java/org/apache/beam/sdk/extensions/sql/impl/schema/BeamTableUtils.java @@ -167,6 +167,8 @@ public static Object autoCastField(Schema.Field field, @Nullable Object rawObj) return Float.valueOf(raw); case DOUBLE: return Double.valueOf(raw); + case DECIMAL: + return new BigDecimal(raw); default: throw new UnsupportedOperationException( String.format("Column type %s is not supported yet!", type)); diff --git a/sdks/java/extensions/sql/src/main/java/org/apache/beam/sdk/extensions/sql/impl/transform/BeamBuiltinAggregations.java b/sdks/java/extensions/sql/src/main/java/org/apache/beam/sdk/extensions/sql/impl/transform/BeamBuiltinAggregations.java index fa2b6fe4adcd..f79527e4da63 100644 --- a/sdks/java/extensions/sql/src/main/java/org/apache/beam/sdk/extensions/sql/impl/transform/BeamBuiltinAggregations.java +++ b/sdks/java/extensions/sql/src/main/java/org/apache/beam/sdk/extensions/sql/impl/transform/BeamBuiltinAggregations.java @@ -526,7 +526,10 @@ public BigDecimal toBigDecimal(Byte record) { static class FloatAvg extends Avg { @Override public @Nullable Float extractOutput(KV accumulator) { - return accumulator.getKey() == 0 ? null : prepareOutput(accumulator).floatValue(); + if (accumulator.getKey() == 0) { + return null; + } + return (float) (accumulator.getValue().doubleValue() / accumulator.getKey()); } @Override @@ -538,7 +541,10 @@ public BigDecimal toBigDecimal(Float record) { static class DoubleAvg extends Avg { @Override public @Nullable Double extractOutput(KV accumulator) { - return accumulator.getKey() == 0 ? null : prepareOutput(accumulator).doubleValue(); + if (accumulator.getKey() == 0) { + return null; + } + return accumulator.getValue().doubleValue() / accumulator.getKey(); } @Override diff --git a/sdks/java/extensions/sql/src/main/java/org/apache/beam/sdk/extensions/sql/impl/transform/BeamBuiltinAnalyticFunctions.java b/sdks/java/extensions/sql/src/main/java/org/apache/beam/sdk/extensions/sql/impl/transform/BeamBuiltinAnalyticFunctions.java index 413093d416d8..b601c45f23e7 100644 --- a/sdks/java/extensions/sql/src/main/java/org/apache/beam/sdk/extensions/sql/impl/transform/BeamBuiltinAnalyticFunctions.java +++ b/sdks/java/extensions/sql/src/main/java/org/apache/beam/sdk/extensions/sql/impl/transform/BeamBuiltinAnalyticFunctions.java @@ -40,11 +40,14 @@ public class BeamBuiltinAnalyticFunctions { // Navigation Functions .put("FIRST_VALUE", typeName -> navigationFirstValue()) .put("LAST_VALUE", typeName -> navigationLastValue()) + // Placeholder entry; BeamWindowRel instantiates NTH_VALUE with the actual offset n. + .put("NTH_VALUE", typeName -> navigationNthValue(1L, false)) // Numbering Functions .put("ROW_NUMBER", typeName -> numberingRowNumber()) .put("DENSE_RANK", typeName -> numberingDenseRank()) .put("RANK", typeName -> numberingRank()) .put("PERCENT_RANK", typeName -> numberingPercentRank()) + .put("CUME_DIST", typeName -> numberingCumeDist()) .build(); public static Combine.CombineFn create(String functionName, Schema.FieldType fieldType) { @@ -66,9 +69,104 @@ public class BeamBuiltinAnalyticFunctions { return new LastValueCombineFn(); } - private static class FirstValueCombineFn extends Combine.CombineFn, T> { + /** + * FIRST_VALUE that optionally skips nulls. With {@code ignoreNulls=true} it returns the first + * non-null value in the frame; otherwise it behaves like {@link #navigationFirstValue()}. + */ + public static Combine.CombineFn navigationFirstValue(boolean ignoreNulls) { + return ignoreNulls ? new FirstValueIgnoreNullsCombineFn() : new FirstValueCombineFn(); + } + + /** + * LAST_VALUE that optionally skips nulls. With {@code ignoreNulls=true} it returns the last + * non-null value in the frame; otherwise it behaves like {@link #navigationLastValue()}. + */ + public static Combine.CombineFn navigationLastValue(boolean ignoreNulls) { + return ignoreNulls ? new LastValueIgnoreNullsCombineFn() : new LastValueCombineFn(); + } + + /** + * NTH_VALUE(expr, n) [IGNORE NULLS]: the value of {@code expr} from the n-th row of the frame + * (frame-absolute, 1-based, counted from the frame start). Returns {@code null} when the frame + * has fewer than n rows. With {@code ignoreNulls=true}, n counts only non-null values. + */ + public static Combine.CombineFn navigationNthValue(long n, boolean ignoreNulls) { + return new NthValueCombineFn<>(n, ignoreNulls); + } + + /** + * Accumulator for FIRST_VALUE / LAST_VALUE that, unlike {@link Optional}, can represent a {@code + * null} value distinct from "no row seen yet". Plain {@code Optional.of(input)} throws an NPE + * when a row's value is null, but FIRST_VALUE / LAST_VALUE (without IGNORE NULLS) must faithfully + * return the first / last row's value even when that value is null. + */ + private static class ValueHolder { + private boolean seen; + private T value; + + void set(T input) { + this.seen = true; + this.value = input; + } + } + + private static class FirstValueCombineFn extends Combine.CombineFn, T> { private FirstValueCombineFn() {} + @Override + public ValueHolder createAccumulator() { + return new ValueHolder<>(); + } + + @Override + public ValueHolder addInput(ValueHolder accumulator, T input) { + if (!accumulator.seen) { + accumulator.set(input); + } + return accumulator; + } + + @Override + public ValueHolder mergeAccumulators(Iterable> accumulators) { + throw new UnsupportedOperationException(); + } + + @Override + public T extractOutput(ValueHolder accumulator) { + return accumulator.seen ? accumulator.value : null; + } + } + + private static class LastValueCombineFn extends Combine.CombineFn, T> { + private LastValueCombineFn() {} + + @Override + public ValueHolder createAccumulator() { + return new ValueHolder<>(); + } + + @Override + public ValueHolder addInput(ValueHolder accumulator, T input) { + accumulator.set(input); + return accumulator; + } + + @Override + public ValueHolder mergeAccumulators(Iterable> accumulators) { + throw new UnsupportedOperationException(); + } + + @Override + public T extractOutput(ValueHolder accumulator) { + return accumulator.seen ? accumulator.value : null; + } + } + + /** FIRST_VALUE(... IGNORE NULLS): the first non-null value in the frame. */ + private static class FirstValueIgnoreNullsCombineFn + extends Combine.CombineFn, T> { + private FirstValueIgnoreNullsCombineFn() {} + @Override public Optional createAccumulator() { return Optional.empty(); @@ -76,11 +174,10 @@ public Optional createAccumulator() { @Override public Optional addInput(Optional accumulator, T input) { - Optional r = accumulator; - if (!accumulator.isPresent()) { - r = Optional.of(input); + if (!accumulator.isPresent() && input != null) { + return Optional.of(input); } - return r; + return accumulator; } @Override @@ -94,8 +191,10 @@ public T extractOutput(Optional accumulator) { } } - private static class LastValueCombineFn extends Combine.CombineFn, T> { - private LastValueCombineFn() {} + /** LAST_VALUE(... IGNORE NULLS): the last non-null value in the frame. */ + private static class LastValueIgnoreNullsCombineFn + extends Combine.CombineFn, T> { + private LastValueIgnoreNullsCombineFn() {} @Override public Optional createAccumulator() { @@ -104,8 +203,10 @@ public Optional createAccumulator() { @Override public Optional addInput(Optional accumulator, T input) { - Optional r = Optional.of(input); - return r; + if (input != null) { + return Optional.of(input); + } + return accumulator; } @Override @@ -119,6 +220,74 @@ public T extractOutput(Optional accumulator) { } } + /** + * NTH_VALUE(expr, n) [IGNORE NULLS]. A position-aware combiner: it receives every frame row's + * value column (via the {@link PositionAwareCombineFn#addInput(Object, Object, Object, Long, + * Long, Long)} overload) and keeps the value seen at the n-th eligible position. Eligibility is + * "every row" by default, or "only non-null values" with IGNORE NULLS. + */ + private static class NthValueCombineFn + extends PositionAwareCombineFn, T> { + + /** Tracks how many eligible rows have been seen and the value captured at position n. */ + private static class NthAccum { + long eligibleSeen = 0L; + boolean captured = false; + T value = null; + } + + private final long n; + private final boolean ignoreNulls; + + private NthValueCombineFn(long n, boolean ignoreNulls) { + this.n = n; + this.ignoreNulls = ignoreNulls; + } + + @Override + public NthAccum createAccumulator() { + return new NthAccum<>(); + } + + @Override + public NthAccum addInput( + NthAccum accumulator, + BigDecimal input, + Long cursorWindow, + Long cursorPartition, + Long countPartition) { + return addInput(accumulator, input, input, cursorWindow, cursorPartition, countPartition); + } + + @Override + @SuppressWarnings("unchecked") + public NthAccum addInput( + NthAccum accumulator, + BigDecimal input, + Object inputValue, + Long cursorWindow, + Long cursorPartition, + Long countPartition) { + if (accumulator.captured) { + return accumulator; + } + if (ignoreNulls && inputValue == null) { + return accumulator; + } + accumulator.eligibleSeen++; + if (accumulator.eligibleSeen == n) { + accumulator.captured = true; + accumulator.value = (T) inputValue; + } + return accumulator; + } + + @Override + public T extractOutput(NthAccum accumulator) { + return accumulator.captured ? accumulator.value : null; + } + } + // Numbering functions public static Combine.CombineFn numberingRowNumber() { return new RowNumberCombineFn(); @@ -136,6 +305,10 @@ public T extractOutput(Optional accumulator) { return new PercentRankCombineFn(); } + public static Combine.CombineFn numberingCumeDist() { + return new CumeDistCombineFn(); + } + public abstract static class PositionAwareCombineFn extends Combine.CombineFn { public abstract AccumT addInput( @@ -145,6 +318,23 @@ public abstract AccumT addInput( Long cursorPartition, Long countPartition); + /** + * Position-aware addInput that, in addition to the ORDER-BY key ({@code input}), receives the + * row's actual VALUE column ({@code inputValue}). Numbering / ranking functions only need the + * ORDER-BY key and ignore {@code inputValue}; navigation functions such as NTH_VALUE need the + * value column at the n-th position within the frame. The default routes to the key-only + * overload so existing combiners are unaffected. + */ + public AccumT addInput( + AccumT accumulator, + InputT input, + Object inputValue, + Long cursorWindow, + Long cursorPartition, + Long countPartition) { + return addInput(accumulator, input, cursorWindow, cursorPartition, countPartition); + } + @Override public AccumT addInput(AccumT mutableAccumulator, InputT input) { throw new UnsupportedOperationException(); @@ -251,6 +441,44 @@ public Long extractOutput(KV accumulator) { } } + /** + * CUME_DIST: the relative rank of the current row, computed as {@code (number of rows preceding + * or peer with the current row) / (total rows in the partition)}. + * + *

Beam evaluates analytic functions over the per-row frame. With the SQL-standard default + * frame for CUME_DIST (RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW), the frame handed to + * this combiner already contains exactly the "preceding or peer" rows, so the numerator is simply + * the number of {@code addInput} calls. The denominator is the partition size, which is supplied + * to every {@code addInput} call as {@code countPartition}. + */ + private static class CumeDistCombineFn + extends PositionAwareCombineFn, Double> { + + @Override + public KV createAccumulator() { + return KV.of(0L, 0L); + } + + @Override + public KV addInput( + KV accumulator, + BigDecimal input, + Long cursorPosition, + Long cursorPartition, + Long countPartition) { + return KV.of(accumulator.getKey() + 1L, countPartition); + } + + @Override + public Double extractOutput(KV accumulator) { + long total = accumulator.getValue(); + if (total <= 0L) { + return 0.0; + } + return accumulator.getKey().doubleValue() / (double) total; + } + } + private static class PercentRankCombineFn extends PositionAwareCombineFn, KV>, Double> { diff --git a/sdks/java/extensions/sql/src/main/java/org/apache/beam/sdk/extensions/sql/impl/transform/agg/VarianceFn.java b/sdks/java/extensions/sql/src/main/java/org/apache/beam/sdk/extensions/sql/impl/transform/agg/VarianceFn.java index 906bac7add52..86f76a28eb7a 100644 --- a/sdks/java/extensions/sql/src/main/java/org/apache/beam/sdk/extensions/sql/impl/transform/agg/VarianceFn.java +++ b/sdks/java/extensions/sql/src/main/java/org/apache/beam/sdk/extensions/sql/impl/transform/agg/VarianceFn.java @@ -30,6 +30,7 @@ import org.apache.beam.sdk.transforms.Combine; import org.apache.beam.sdk.transforms.SerializableFunction; import org.apache.beam.vendor.calcite.v1_40_0.org.apache.calcite.runtime.SqlFunctions; +import org.checkerframework.checker.nullness.qual.Nullable; /** * {@link Combine.CombineFn} for Variance on {@link Number} types. @@ -68,7 +69,8 @@ @SuppressWarnings({ "rawtypes" // TODO(https://github.com/apache/beam/issues/20447) }) -public class VarianceFn extends Combine.CombineFn { +public class VarianceFn + extends Combine.CombineFn { static final MathContext MATH_CTX = new MathContext(10, RoundingMode.HALF_UP); @@ -76,9 +78,6 @@ public class VarianceFn extends Combine.CombineFn decimalConverter; @@ -146,9 +145,12 @@ public Coder getAccumulatorCoder( } @Override - public T extractOutput(VarianceAccumulator accumulator) { + public @Nullable T extractOutput(VarianceAccumulator accumulator) { BigDecimal result = getVariance(accumulator); - if (result != null && isStddev) { + if (result == null) { + return null; + } + if (isStddev) { double doubleVal = result.doubleValue(); if (doubleVal < 0.0) { doubleVal = 0.0; // Clamp negative variance due to numerical instability @@ -162,14 +164,14 @@ public T extractOutput(VarianceAccumulator accumulator) { return decimalConverter.apply(result); } - private BigDecimal getVariance(VarianceAccumulator variance) { + private @Nullable BigDecimal getVariance(VarianceAccumulator variance) { BigDecimal adjustedCount = this.isSample ? variance.count().subtract(BigDecimal.ONE) : variance.count(); - // Avoid ArithmeticException: Division is undefined when adjustedCount is 0 - if (adjustedCount.compareTo(BigDecimal.ZERO) == 0) { - return BigDecimal.ZERO; + // Sample variance of <= 1 element or population variance of 0 elements is undefined (NULL). + if (adjustedCount.compareTo(BigDecimal.ZERO) <= 0) { + return null; } return variance.variance().divide(adjustedCount, MATH_CTX); diff --git a/sdks/java/extensions/sql/src/test/java/org/apache/beam/sdk/extensions/sql/BeamAnalyticFunctionsTest.java b/sdks/java/extensions/sql/src/test/java/org/apache/beam/sdk/extensions/sql/BeamAnalyticFunctionsTest.java index d66fdb58d4af..a3719267030e 100644 --- a/sdks/java/extensions/sql/src/test/java/org/apache/beam/sdk/extensions/sql/BeamAnalyticFunctionsTest.java +++ b/sdks/java/extensions/sql/src/test/java/org/apache/beam/sdk/extensions/sql/BeamAnalyticFunctionsTest.java @@ -588,4 +588,116 @@ public void testPercentRankFunction() throws Exception { pipeline.run(); } + + @Test + public void testCumeDistFunction() throws Exception { + pipeline.enableAbandonedNodeEnforcement(false); + PCollection inputRows = inputData2(); + String sql = "SELECT x, CUME_DIST() over (ORDER BY x ) as agg FROM PCOLLECTION"; + PCollection result = inputRows.apply("sql", SqlTransform.query(sql)); + + Schema overResultSchema = Schema.builder().addInt32Field("x").addDoubleField("agg").build(); + + List overResult = + TestUtils.RowsBuilder.of(overResultSchema) + .addRows( + 1, 1.0 / 7.0, + 2, 3.0 / 7.0, + 2, 3.0 / 7.0, + 5, 4.0 / 7.0, + 8, 5.0 / 7.0, + 10, 7.0 / 7.0, + 10, 7.0 / 7.0) + .getRows(); + + PAssert.that(result).containsInAnyOrder(overResult); + + pipeline.run(); + } + + @Test + public void testCountStarOverWindow() throws Exception { + pipeline.enableAbandonedNodeEnforcement(false); + PCollection inputRows = inputData2(); + String sql = "SELECT x, COUNT(*) over () as agg FROM PCOLLECTION"; + PCollection result = inputRows.apply("sql", SqlTransform.query(sql)); + + Schema overResultSchema = Schema.builder().addInt32Field("x").addInt64Field("agg").build(); + + List overResult = + TestUtils.RowsBuilder.of(overResultSchema) + .addRows( + 1, 7L, + 2, 7L, + 2, 7L, + 5, 7L, + 8, 7L, + 10, 7L, + 10, 7L) + .getRows(); + + PAssert.that(result).containsInAnyOrder(overResult); + + pipeline.run(); + } + + /** + * Table with a TIMESTAMP ordering key. Used to exercise RANGE framing over a temporal ORDER BY + * column (previously a ClassCastException because the order-by value was cast to {@code Number}). + */ + private PCollection inputDataTimestamp() { + Schema schema = Schema.builder().addDateTimeField("ts").addInt32Field("v").build(); + return pipeline + .apply( + Create.of( + TestUtils.rowsBuilderOf(schema) + .addRows( + org.apache.beam.sdk.extensions.sql.utils.DateTimeUtils + .parseTimestampWithUTCTimeZone("2020-01-01 00:00:00"), + 10, + org.apache.beam.sdk.extensions.sql.utils.DateTimeUtils + .parseTimestampWithUTCTimeZone("2020-01-02 00:00:00"), + 20, + org.apache.beam.sdk.extensions.sql.utils.DateTimeUtils + .parseTimestampWithUTCTimeZone("2020-01-02 00:00:00"), + 30, + org.apache.beam.sdk.extensions.sql.utils.DateTimeUtils + .parseTimestampWithUTCTimeZone("2020-01-03 00:00:00"), + 40) + .getRows())) + .setRowSchema(schema); + } + + /** + * RANGE framing over a TIMESTAMP ORDER BY key. A cumulative RANGE sum groups rows that share the + * same ordering value into one peer group: the two rows at 2020-01-02 both see 10+20+30=60. + */ + @Test + public void testOverRangeTimestampOrderKey() throws Exception { + pipeline.enableAbandonedNodeEnforcement(false); + PCollection inputRows = inputDataTimestamp(); + String sql = + "SELECT v, sum(v) over " + + "(" + + "ORDER BY ts " + + "RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW" + + ")" + + " as total FROM PCOLLECTION"; + PCollection result = inputRows.apply("sql", SqlTransform.query(sql)); + + Schema overResultSchema = Schema.builder().addInt32Field("v").addInt32Field("total").build(); + + List overResult = + TestUtils.RowsBuilder.of(overResultSchema) + .addRows( + 10, 10, + 20, 60, + 30, 60, + 40, 100) + .getRows(); + + PAssert.that(result).containsInAnyOrder(overResult); + + pipeline.run(); + } } diff --git a/sdks/java/extensions/sql/src/test/java/org/apache/beam/sdk/extensions/sql/BeamSqlDslJoinTest.java b/sdks/java/extensions/sql/src/test/java/org/apache/beam/sdk/extensions/sql/BeamSqlDslJoinTest.java index b210bb0258f5..8bae1107b616 100644 --- a/sdks/java/extensions/sql/src/test/java/org/apache/beam/sdk/extensions/sql/BeamSqlDslJoinTest.java +++ b/sdks/java/extensions/sql/src/test/java/org/apache/beam/sdk/extensions/sql/BeamSqlDslJoinTest.java @@ -146,12 +146,17 @@ public void testException_nonEqualJoin() throws Exception { pipeline.run(); } - @Test(expected = UnsupportedOperationException.class) - public void testException_crossJoin() throws Exception { + @Test + public void testCrossJoin() throws Exception { String sql = "SELECT * " + "FROM ORDER_DETAILS1 o1, ORDER_DETAILS2 o2"; - pipeline.enableAbandonedNodeEnforcement(false); - queryFromOrderTables(sql); + PAssert.that(queryFromOrderTables(sql)) + .containsInAnyOrder( + TestUtils.RowsBuilder.of(RESULT_ROW_TYPE) + .addRows( + 1, 2, 3, 1, 2, 3, 1, 2, 3, 2, 3, 3, 1, 2, 3, 3, 4, 5, 2, 3, 3, 1, 2, 3, 2, 3, 3, + 2, 3, 3, 2, 3, 3, 3, 4, 5, 3, 4, 5, 1, 2, 3, 3, 4, 5, 2, 3, 3, 3, 4, 5, 3, 4, 5) + .getRows()); pipeline.run(); } diff --git a/sdks/java/extensions/sql/src/test/java/org/apache/beam/sdk/extensions/sql/impl/rel/BeamCoGBKJoinRelBoundedVsBoundedTest.java b/sdks/java/extensions/sql/src/test/java/org/apache/beam/sdk/extensions/sql/impl/rel/BeamCoGBKJoinRelBoundedVsBoundedTest.java index 9a0977e9018c..27c9e14d8748 100644 --- a/sdks/java/extensions/sql/src/test/java/org/apache/beam/sdk/extensions/sql/impl/rel/BeamCoGBKJoinRelBoundedVsBoundedTest.java +++ b/sdks/java/extensions/sql/src/test/java/org/apache/beam/sdk/extensions/sql/impl/rel/BeamCoGBKJoinRelBoundedVsBoundedTest.java @@ -356,7 +356,7 @@ public void testException_join_condition1() throws Exception { } @Test - public void testException_join_condition2() throws Exception { + public void testJoinWithEquiAndNonEquiCondition() throws Exception { String sql = "SELECT * " + "FROM ORDER_DETAILS1 o1" @@ -364,9 +364,21 @@ public void testException_join_condition2() throws Exception { + " on " + " o1.order_id = o2.site_id AND o1.price > o2.site_id"; - thrown.expect(UnsupportedOperationException.class); - thrown.expectMessage(StringContains.containsString("Non equi-join")); - compilePipeline(sql, pipeline); + PCollection rows = compilePipeline(sql, pipeline); + PAssert.that(rows) + .containsInAnyOrder( + TestUtils.RowsBuilder.of( + Schema.builder() + .addField("order_id", Schema.FieldType.INT32) + .addField("site_id", Schema.FieldType.INT32) + .addField("price", Schema.FieldType.INT32) + .addField("order_id0", Schema.FieldType.INT32) + .addField("site_id0", Schema.FieldType.INT32) + .addField("price0", Schema.FieldType.INT32) + .build()) + .addRows(2, 3, 3, 1, 2, 3) + .addRows(3, 4, 5, 2, 3, 3) + .getRows()); pipeline.run(); } @@ -402,12 +414,26 @@ public void testException_join_condition4() throws Exception { pipeline.run(); } - @Test(expected = UnsupportedOperationException.class) - public void testException_crossJoin() throws Exception { + @Test + public void testCrossJoin() throws Exception { String sql = "SELECT * " + "FROM ORDER_DETAILS1 o1, ORDER_DETAILS2 o2"; - pipeline.enableAbandonedNodeEnforcement(false); - compilePipeline(sql, pipeline); + PCollection rows = compilePipeline(sql, pipeline); + PAssert.that(rows) + .containsInAnyOrder( + TestUtils.RowsBuilder.of( + Schema.builder() + .addField("order_id", Schema.FieldType.INT32) + .addField("site_id", Schema.FieldType.INT32) + .addField("price", Schema.FieldType.INT32) + .addField("order_id0", Schema.FieldType.INT32) + .addField("site_id0", Schema.FieldType.INT32) + .addField("price0", Schema.FieldType.INT32) + .build()) + .addRows( + 1, 2, 3, 1, 2, 3, 1, 2, 3, 2, 3, 3, 1, 2, 3, 3, 4, 5, 2, 3, 3, 1, 2, 3, 2, 3, 3, + 2, 3, 3, 2, 3, 3, 3, 4, 5, 3, 4, 5, 1, 2, 3, 3, 4, 5, 2, 3, 3, 3, 4, 5, 3, 4, 5) + .getRows()); pipeline.run(); } } diff --git a/sdks/java/testing/tpcds/src/main/java/org/apache/beam/sdk/tpcds/QueryReader.java b/sdks/java/testing/tpcds/src/main/java/org/apache/beam/sdk/tpcds/QueryReader.java index f6bb0d500920..14a951d76cbf 100644 --- a/sdks/java/testing/tpcds/src/main/java/org/apache/beam/sdk/tpcds/QueryReader.java +++ b/sdks/java/testing/tpcds/src/main/java/org/apache/beam/sdk/tpcds/QueryReader.java @@ -19,6 +19,7 @@ import java.nio.charset.StandardCharsets; import java.util.Set; +import org.apache.beam.vendor.calcite.v1_40_0.org.apache.calcite.avatica.util.Quoting; import org.apache.beam.vendor.calcite.v1_40_0.org.apache.calcite.sql.SqlNode; import org.apache.beam.vendor.calcite.v1_40_0.org.apache.calcite.sql.parser.SqlParseException; import org.apache.beam.vendor.calcite.v1_40_0.org.apache.calcite.sql.parser.SqlParser; @@ -50,7 +51,8 @@ public static String readQuery(String queryFileName) throws Exception { * @throws SqlParseException */ public static Set getQueryIdentifiers(String queryString) throws SqlParseException { - SqlParser parser = SqlParser.create(queryString); + SqlParser parser = + SqlParser.create(queryString, SqlParser.config().withQuoting(Quoting.BACK_TICK)); SqlNode parsedQuery = parser.parseQuery(); SqlTransformRunner.SqlIdentifierVisitor sqlVisitor = new SqlTransformRunner.SqlIdentifierVisitor(); diff --git a/sdks/java/testing/tpcds/src/main/java/org/apache/beam/sdk/tpcds/TableSchemaJSONLoader.java b/sdks/java/testing/tpcds/src/main/java/org/apache/beam/sdk/tpcds/TableSchemaJSONLoader.java index 0d683c6b66da..cd5f8ad79883 100644 --- a/sdks/java/testing/tpcds/src/main/java/org/apache/beam/sdk/tpcds/TableSchemaJSONLoader.java +++ b/sdks/java/testing/tpcds/src/main/java/org/apache/beam/sdk/tpcds/TableSchemaJSONLoader.java @@ -78,9 +78,7 @@ public static String parseTableSchema(String tableName) throws Exception { // Use long type to represent int, prevent overflow schemaStringBuilder.append("bigint"); } else if (value.contains("decimal")) { - // Currently Beam SQL doesn't handle "decimal" type properly, use "double" to replace it - // for now. - schemaStringBuilder.append("double"); + schemaStringBuilder.append("decimal"); } else { // Currently Beam SQL doesn't handle "date" type properly, use "varchar" replace it for // now. diff --git a/sdks/java/testing/tpcds/src/main/java/org/apache/beam/sdk/tpcds/TpcdsOptions.java b/sdks/java/testing/tpcds/src/main/java/org/apache/beam/sdk/tpcds/TpcdsOptions.java index 93cbc94f30a0..8394acd70dbe 100644 --- a/sdks/java/testing/tpcds/src/main/java/org/apache/beam/sdk/tpcds/TpcdsOptions.java +++ b/sdks/java/testing/tpcds/src/main/java/org/apache/beam/sdk/tpcds/TpcdsOptions.java @@ -36,8 +36,9 @@ public interface TpcdsOptions extends ApplicationNameOptions, GcpOptions, BeamSq void setDataSize(String dataSize); // Set the return type to be String since reading from the command line (user input will be like - // "1,2,55" which represent TPC-DS query1, query3, query55) - @Description("The queries numbers, read user input as string, numbers separated by commas") + // "1,2,55" which represent TPC-DS query1, query3, query55, or "supported" / "all") + @Description( + "The queries numbers, read user input as string, numbers separated by commas, or 'supported' / 'all'") String getQueries(); void setQueries(String queries); diff --git a/sdks/java/testing/tpcds/src/main/java/org/apache/beam/sdk/tpcds/TpcdsParametersReader.java b/sdks/java/testing/tpcds/src/main/java/org/apache/beam/sdk/tpcds/TpcdsParametersReader.java index 60072ea916c0..a2a451b554e8 100644 --- a/sdks/java/testing/tpcds/src/main/java/org/apache/beam/sdk/tpcds/TpcdsParametersReader.java +++ b/sdks/java/testing/tpcds/src/main/java/org/apache/beam/sdk/tpcds/TpcdsParametersReader.java @@ -30,13 +30,22 @@ public class TpcdsParametersReader { /** The data sizes that have been supported. */ private static final Set supportedDataSizes = - Stream.of("1G", "1GB", "10G", "10GB", "100G", "100GB") + Stream.of("10M", "10MB", "1G", "1GB", "10G", "10GB", "100G", "100GB") .collect(Collectors.toCollection(HashSet::new)); private static final String QUERY_PREFIX = "query"; public static final List ALL_QUERY_NAMES = getAllQueryNames(); + public static final List SUPPORTED_QUERY_NAMES = + Stream.of( + "1", "3", "4", "7", "9", "10", "11", "12", "15", "17", "19", "20", "21", "26", "31", + "32", "33", "34", "35", "38", "40", "42", "43", "45", "46", "49", "50", "52", "53", + "55", "56", "59", "60", "61", "63", "65", "66", "68", "69", "71", "73", "74", "75", + "76", "79", "82", "84", "87", "88", "89", "92", "93", "96", "97", "99") + .map(s -> QUERY_PREFIX + s) + .collect(Collectors.toList()); + /** * Get and check dataSize entered by user. This dataSize has to have been supported. * @@ -55,9 +64,10 @@ public static String getAndCheckDataSize(TpcdsOptions tpcdsOptions) throws Excep } /** - * Get and check queries entered by user. This has to be a string of numbers separated by commas - * or "all" which means run all 99 queries. All query numbers have to be between 1 and 99. Some - * queries (14, 23, 24 and 39) may contain suffixes "a" or "b", e.g. "14a" and "14b". + * Get and check queries entered by user. This has to be a string of numbers separated by commas, + * "supported" which means run all supported queries, or "all" which means run all 99 queries. All + * query numbers have to be between 1 and 99. Some queries (14, 23, 24 and 39) may contain + * suffixes "a" or "b", e.g. "14a" and "14b". * * @param tpcdsOptions TpcdsOptions object constructed from user input * @return An array of query names, for example "1,2,14b" will be output as @@ -67,6 +77,8 @@ public static String[] getAndCheckQueryNames(TpcdsOptions tpcdsOptions) { if (tpcdsOptions.getQueries().equalsIgnoreCase("all")) { // All TPC-DS queries need to be executed. return ALL_QUERY_NAMES.toArray(new String[0]); + } else if (tpcdsOptions.getQueries().equalsIgnoreCase("supported")) { + return SUPPORTED_QUERY_NAMES.toArray(new String[0]); } else { List queries = new ArrayList<>(); Arrays.stream(tpcdsOptions.getQueries().split("[\\s,]+", -1)) diff --git a/sdks/java/testing/tpcds/src/main/java/org/apache/beam/sdk/tpcds/TpcdsSchemas.java b/sdks/java/testing/tpcds/src/main/java/org/apache/beam/sdk/tpcds/TpcdsSchemas.java index 763e8164dee0..196e27e039c7 100644 --- a/sdks/java/testing/tpcds/src/main/java/org/apache/beam/sdk/tpcds/TpcdsSchemas.java +++ b/sdks/java/testing/tpcds/src/main/java/org/apache/beam/sdk/tpcds/TpcdsSchemas.java @@ -58,8 +58,8 @@ public static Map getTpcdsSchemas() throws Exception { Schema.FieldType fieldType; if (type.equals("bigint")) { fieldType = Schema.FieldType.INT64; - } else if (type.equals("double")) { - fieldType = Schema.FieldType.DOUBLE; + } else if (type.equals("decimal")) { + fieldType = Schema.FieldType.DECIMAL; } else { fieldType = Schema.FieldType.STRING; } @@ -235,8 +235,8 @@ public static Schema getWebSiteSchema() { .addNullableField("cc_state", Schema.FieldType.STRING) .addNullableField("cc_zip", Schema.FieldType.STRING) .addNullableField("cc_country", Schema.FieldType.STRING) - .addNullableField("cc_gmt_offset", Schema.FieldType.DOUBLE) - .addNullableField("cc_tax_percentage", Schema.FieldType.DOUBLE) + .addNullableField("cc_gmt_offset", Schema.FieldType.DECIMAL) + .addNullableField("cc_tax_percentage", Schema.FieldType.DECIMAL) .build(); private static final Schema catalogPageSchema = @@ -272,15 +272,15 @@ public static Schema getWebSiteSchema() { .addNullableField("cr_reason_sk", Schema.FieldType.INT64) .addField("cr_order_number", Schema.FieldType.INT64) .addNullableField("cr_return_quantity", Schema.FieldType.INT64) - .addNullableField("cr_return_amount", Schema.FieldType.DOUBLE) - .addNullableField("cr_return_tax", Schema.FieldType.DOUBLE) - .addNullableField("cr_return_amt_inc_tax", Schema.FieldType.DOUBLE) - .addNullableField("cr_fee", Schema.FieldType.DOUBLE) - .addNullableField("cr_return_ship_cost", Schema.FieldType.DOUBLE) - .addNullableField("cr_refunded_cash", Schema.FieldType.DOUBLE) - .addNullableField("cr_reversed_charge", Schema.FieldType.DOUBLE) - .addNullableField("cr_store_credit", Schema.FieldType.DOUBLE) - .addNullableField("cr_net_loss", Schema.FieldType.DOUBLE) + .addNullableField("cr_return_amount", Schema.FieldType.DECIMAL) + .addNullableField("cr_return_tax", Schema.FieldType.DECIMAL) + .addNullableField("cr_return_amt_inc_tax", Schema.FieldType.DECIMAL) + .addNullableField("cr_fee", Schema.FieldType.DECIMAL) + .addNullableField("cr_return_ship_cost", Schema.FieldType.DECIMAL) + .addNullableField("cr_refunded_cash", Schema.FieldType.DECIMAL) + .addNullableField("cr_reversed_charge", Schema.FieldType.DECIMAL) + .addNullableField("cr_store_credit", Schema.FieldType.DECIMAL) + .addNullableField("cr_net_loss", Schema.FieldType.DECIMAL) .build(); private static final Schema catalogSalesSchema = @@ -303,22 +303,22 @@ public static Schema getWebSiteSchema() { .addField("cs_item_sk", Schema.FieldType.INT64) .addNullableField("cs_promo_sk", Schema.FieldType.INT64) .addField("cs_order_number", Schema.FieldType.INT64) - .addNullableField("cs_quantity", Schema.FieldType.DOUBLE) - .addNullableField("cs_wholesale_cost", Schema.FieldType.DOUBLE) - .addNullableField("cs_list_price", Schema.FieldType.DOUBLE) - .addNullableField("cs_sales_price", Schema.FieldType.DOUBLE) - .addNullableField("cs_ext_discount_amt", Schema.FieldType.DOUBLE) - .addNullableField("cs_ext_sales_price", Schema.FieldType.DOUBLE) - .addNullableField("cs_ext_wholesale_cost", Schema.FieldType.DOUBLE) - .addNullableField("cs_ext_list_price", Schema.FieldType.DOUBLE) - .addNullableField("cs_ext_tax", Schema.FieldType.DOUBLE) - .addNullableField("cs_coupon_amt", Schema.FieldType.DOUBLE) - .addNullableField("cs_ext_ship_cost", Schema.FieldType.DOUBLE) - .addNullableField("cs_net_paid", Schema.FieldType.DOUBLE) - .addNullableField("cs_net_paid_inc_tax", Schema.FieldType.DOUBLE) - .addNullableField("cs_net_paid_inc_ship", Schema.FieldType.DOUBLE) - .addNullableField("cs_net_paid_inc_ship_tax", Schema.FieldType.DOUBLE) - .addNullableField("cs_net_profit", Schema.FieldType.DOUBLE) + .addNullableField("cs_quantity", Schema.FieldType.INT64) + .addNullableField("cs_wholesale_cost", Schema.FieldType.DECIMAL) + .addNullableField("cs_list_price", Schema.FieldType.DECIMAL) + .addNullableField("cs_sales_price", Schema.FieldType.DECIMAL) + .addNullableField("cs_ext_discount_amt", Schema.FieldType.DECIMAL) + .addNullableField("cs_ext_sales_price", Schema.FieldType.DECIMAL) + .addNullableField("cs_ext_wholesale_cost", Schema.FieldType.DECIMAL) + .addNullableField("cs_ext_list_price", Schema.FieldType.DECIMAL) + .addNullableField("cs_ext_tax", Schema.FieldType.DECIMAL) + .addNullableField("cs_coupon_amt", Schema.FieldType.DECIMAL) + .addNullableField("cs_ext_ship_cost", Schema.FieldType.DECIMAL) + .addNullableField("cs_net_paid", Schema.FieldType.DECIMAL) + .addNullableField("cs_net_paid_inc_tax", Schema.FieldType.DECIMAL) + .addNullableField("cs_net_paid_inc_ship", Schema.FieldType.DECIMAL) + .addNullableField("cs_net_paid_inc_ship_tax", Schema.FieldType.DECIMAL) + .addNullableField("cs_net_profit", Schema.FieldType.DECIMAL) .build(); private static final Schema customerSchema = @@ -356,7 +356,7 @@ public static Schema getWebSiteSchema() { .addNullableField("ca_state", Schema.FieldType.STRING) .addNullableField("ca_zip", Schema.FieldType.STRING) .addNullableField("ca_country", Schema.FieldType.STRING) - .addNullableField("ca_gmt_offset", Schema.FieldType.DOUBLE) + .addNullableField("ca_gmt_offset", Schema.FieldType.DECIMAL) .addNullableField("ca_location_type", Schema.FieldType.STRING) .build(); @@ -436,8 +436,8 @@ public static Schema getWebSiteSchema() { .addNullableField("i_rec_start_date", Schema.FieldType.STRING) .addNullableField("i_rec_end_date", Schema.FieldType.STRING) .addNullableField("i_item_desc", Schema.FieldType.STRING) - .addNullableField("i_current_price", Schema.FieldType.DOUBLE) - .addNullableField("i_wholesale_cost", Schema.FieldType.DOUBLE) + .addNullableField("i_current_price", Schema.FieldType.DECIMAL) + .addNullableField("i_wholesale_cost", Schema.FieldType.DECIMAL) .addNullableField("i_brand_id", Schema.FieldType.INT64) .addNullableField("i_brand", Schema.FieldType.STRING) .addNullableField("i_class_id", Schema.FieldType.INT64) @@ -462,7 +462,7 @@ public static Schema getWebSiteSchema() { .addNullableField("p_start_date_sk", Schema.FieldType.INT64) .addNullableField("p_end_date_sk", Schema.FieldType.INT64) .addNullableField("p_item_sk", Schema.FieldType.INT64) - .addNullableField("p_cost", Schema.FieldType.DOUBLE) + .addNullableField("p_cost", Schema.FieldType.DECIMAL) .addNullableField("p_response_target", Schema.FieldType.INT64) .addNullableField("p_promo_name", Schema.FieldType.STRING) .addNullableField("p_channel_dmail", Schema.FieldType.STRING) @@ -524,8 +524,8 @@ public static Schema getWebSiteSchema() { .addNullableField("s_state", Schema.FieldType.STRING) .addNullableField("s_zip", Schema.FieldType.STRING) .addNullableField("s_country", Schema.FieldType.STRING) - .addNullableField("s_gmt_offset", Schema.FieldType.DOUBLE) - .addNullableField("s_tax_percentage", Schema.FieldType.DOUBLE) + .addNullableField("s_gmt_offset", Schema.FieldType.DECIMAL) + .addNullableField("s_tax_percentage", Schema.FieldType.DECIMAL) .build(); private static final Schema storeReturnsSchema = @@ -541,15 +541,15 @@ public static Schema getWebSiteSchema() { .addNullableField("sr_reason_sk", Schema.FieldType.INT64) .addField("sr_ticket_number", Schema.FieldType.INT64) .addNullableField("sr_return_quantity", Schema.FieldType.INT64) - .addNullableField("sr_return_amt", Schema.FieldType.DOUBLE) - .addNullableField("sr_return_tax", Schema.FieldType.DOUBLE) - .addNullableField("sr_return_amt_inc_tax", Schema.FieldType.DOUBLE) - .addNullableField("sr_fee", Schema.FieldType.DOUBLE) - .addNullableField("sr_return_ship_cost", Schema.FieldType.DOUBLE) - .addNullableField("sr_refunded_cash", Schema.FieldType.DOUBLE) - .addNullableField("sr_reversed_charge", Schema.FieldType.DOUBLE) - .addNullableField("sr_store_credit", Schema.FieldType.DOUBLE) - .addNullableField("sr_net_loss", Schema.FieldType.DOUBLE) + .addNullableField("sr_return_amt", Schema.FieldType.DECIMAL) + .addNullableField("sr_return_tax", Schema.FieldType.DECIMAL) + .addNullableField("sr_return_amt_inc_tax", Schema.FieldType.DECIMAL) + .addNullableField("sr_fee", Schema.FieldType.DECIMAL) + .addNullableField("sr_return_ship_cost", Schema.FieldType.DECIMAL) + .addNullableField("sr_refunded_cash", Schema.FieldType.DECIMAL) + .addNullableField("sr_reversed_charge", Schema.FieldType.DECIMAL) + .addNullableField("sr_store_credit", Schema.FieldType.DECIMAL) + .addNullableField("sr_net_loss", Schema.FieldType.DECIMAL) .build(); private static final Schema storeSalesSchema = @@ -565,18 +565,18 @@ public static Schema getWebSiteSchema() { .addNullableField("ss_promo_sk", Schema.FieldType.INT64) .addField("ss_ticket_number", Schema.FieldType.INT64) .addNullableField("ss_quantity", Schema.FieldType.INT64) - .addNullableField("ss_wholesale_cost", Schema.FieldType.DOUBLE) - .addNullableField("ss_list_price", Schema.FieldType.DOUBLE) - .addNullableField("ss_sales_price", Schema.FieldType.DOUBLE) - .addNullableField("ss_ext_discount_amt", Schema.FieldType.DOUBLE) - .addNullableField("ss_ext_sales_price", Schema.FieldType.DOUBLE) - .addNullableField("ss_ext_wholesale_cost", Schema.FieldType.DOUBLE) - .addNullableField("ss_ext_list_price", Schema.FieldType.DOUBLE) - .addNullableField("ss_ext_tax", Schema.FieldType.DOUBLE) - .addNullableField("ss_coupon_amt", Schema.FieldType.DOUBLE) - .addNullableField("ss_net_paid", Schema.FieldType.DOUBLE) - .addNullableField("ss_net_paid_inc_tax", Schema.FieldType.DOUBLE) - .addNullableField("ss_net_profit", Schema.FieldType.DOUBLE) + .addNullableField("ss_wholesale_cost", Schema.FieldType.DECIMAL) + .addNullableField("ss_list_price", Schema.FieldType.DECIMAL) + .addNullableField("ss_sales_price", Schema.FieldType.DECIMAL) + .addNullableField("ss_ext_discount_amt", Schema.FieldType.DECIMAL) + .addNullableField("ss_ext_sales_price", Schema.FieldType.DECIMAL) + .addNullableField("ss_ext_wholesale_cost", Schema.FieldType.DECIMAL) + .addNullableField("ss_ext_list_price", Schema.FieldType.DECIMAL) + .addNullableField("ss_ext_tax", Schema.FieldType.DECIMAL) + .addNullableField("ss_coupon_amt", Schema.FieldType.DECIMAL) + .addNullableField("ss_net_paid", Schema.FieldType.DECIMAL) + .addNullableField("ss_net_paid_inc_tax", Schema.FieldType.DECIMAL) + .addNullableField("ss_net_profit", Schema.FieldType.DECIMAL) .build(); private static final Schema timeDimSchema = @@ -608,7 +608,7 @@ public static Schema getWebSiteSchema() { .addNullableField("w_state", Schema.FieldType.STRING) .addNullableField("w_zip", Schema.FieldType.STRING) .addNullableField("w_country", Schema.FieldType.STRING) - .addNullableField("w_gmt_offset", Schema.FieldType.DOUBLE) + .addNullableField("w_gmt_offset", Schema.FieldType.DECIMAL) .build(); private static final Schema webPageSchema = @@ -646,15 +646,15 @@ public static Schema getWebSiteSchema() { .addNullableField("wr_reason_sk", Schema.FieldType.INT64) .addField("wr_order_number", Schema.FieldType.INT64) .addNullableField("wr_return_quantity", Schema.FieldType.INT64) - .addNullableField("wr_return_amt", Schema.FieldType.DOUBLE) - .addNullableField("wr_return_tax", Schema.FieldType.DOUBLE) - .addNullableField("wr_return_amt_inc_tax", Schema.FieldType.DOUBLE) - .addNullableField("wr_fee", Schema.FieldType.DOUBLE) - .addNullableField("wr_return_ship_cost", Schema.FieldType.DOUBLE) - .addNullableField("wr_refunded_cash", Schema.FieldType.DOUBLE) - .addNullableField("wr_reversed_charge", Schema.FieldType.DOUBLE) - .addNullableField("wr_account_credit", Schema.FieldType.DOUBLE) - .addNullableField("wr_net_loss", Schema.FieldType.DOUBLE) + .addNullableField("wr_return_amt", Schema.FieldType.DECIMAL) + .addNullableField("wr_return_tax", Schema.FieldType.DECIMAL) + .addNullableField("wr_return_amt_inc_tax", Schema.FieldType.DECIMAL) + .addNullableField("wr_fee", Schema.FieldType.DECIMAL) + .addNullableField("wr_return_ship_cost", Schema.FieldType.DECIMAL) + .addNullableField("wr_refunded_cash", Schema.FieldType.DECIMAL) + .addNullableField("wr_reversed_charge", Schema.FieldType.DECIMAL) + .addNullableField("wr_account_credit", Schema.FieldType.DECIMAL) + .addNullableField("wr_net_loss", Schema.FieldType.DECIMAL) .build(); private static final Schema webSalesSchema = @@ -678,21 +678,21 @@ public static Schema getWebSiteSchema() { .addNullableField("ws_promo_sk", Schema.FieldType.INT32) .addField("ws_order_number", Schema.FieldType.INT64) .addNullableField("ws_quantity", Schema.FieldType.INT32) - .addNullableField("ws_wholesale_cost", Schema.FieldType.DOUBLE) - .addNullableField("ws_list_price", Schema.FieldType.DOUBLE) - .addNullableField("ws_sales_price", Schema.FieldType.DOUBLE) - .addNullableField("ws_ext_discount_amt", Schema.FieldType.DOUBLE) - .addNullableField("ws_ext_sales_price", Schema.FieldType.DOUBLE) - .addNullableField("ws_ext_wholesale_cost", Schema.FieldType.DOUBLE) - .addNullableField("ws_ext_list_price", Schema.FieldType.DOUBLE) - .addNullableField("ws_ext_tax", Schema.FieldType.DOUBLE) - .addNullableField("ws_coupon_amt", Schema.FieldType.DOUBLE) - .addNullableField("ws_ext_ship_cost", Schema.FieldType.DOUBLE) - .addNullableField("ws_net_paid", Schema.FieldType.DOUBLE) - .addNullableField("ws_net_paid_inc_tax", Schema.FieldType.DOUBLE) - .addNullableField("ws_net_paid_inc_ship", Schema.FieldType.DOUBLE) - .addNullableField("ws_net_paid_inc_ship_tax", Schema.FieldType.DOUBLE) - .addNullableField("ws_net_profit", Schema.FieldType.DOUBLE) + .addNullableField("ws_wholesale_cost", Schema.FieldType.DECIMAL) + .addNullableField("ws_list_price", Schema.FieldType.DECIMAL) + .addNullableField("ws_sales_price", Schema.FieldType.DECIMAL) + .addNullableField("ws_ext_discount_amt", Schema.FieldType.DECIMAL) + .addNullableField("ws_ext_sales_price", Schema.FieldType.DECIMAL) + .addNullableField("ws_ext_wholesale_cost", Schema.FieldType.DECIMAL) + .addNullableField("ws_ext_list_price", Schema.FieldType.DECIMAL) + .addNullableField("ws_ext_tax", Schema.FieldType.DECIMAL) + .addNullableField("ws_coupon_amt", Schema.FieldType.DECIMAL) + .addNullableField("ws_ext_ship_cost", Schema.FieldType.DECIMAL) + .addNullableField("ws_net_paid", Schema.FieldType.DECIMAL) + .addNullableField("ws_net_paid_inc_tax", Schema.FieldType.DECIMAL) + .addNullableField("ws_net_paid_inc_ship", Schema.FieldType.DECIMAL) + .addNullableField("ws_net_paid_inc_ship_tax", Schema.FieldType.DECIMAL) + .addNullableField("ws_net_profit", Schema.FieldType.DECIMAL) .build(); private static final Schema webSiteSchema = @@ -721,7 +721,7 @@ public static Schema getWebSiteSchema() { .addNullableField("web_state", Schema.FieldType.STRING) .addNullableField("web_zip", Schema.FieldType.STRING) .addNullableField("web_country", Schema.FieldType.STRING) - .addNullableField("web_gmt_offset", Schema.FieldType.DOUBLE) - .addNullableField("web_tax_percentage", Schema.FieldType.DOUBLE) + .addNullableField("web_gmt_offset", Schema.FieldType.DECIMAL) + .addNullableField("web_tax_percentage", Schema.FieldType.DECIMAL) .build(); } diff --git a/sdks/java/testing/tpcds/src/test/java/org/apache/beam/sdk/tpcds/TableSchemaJSONLoaderTest.java b/sdks/java/testing/tpcds/src/test/java/org/apache/beam/sdk/tpcds/TableSchemaJSONLoaderTest.java index 3eb84d32c57c..1fd0422c8146 100644 --- a/sdks/java/testing/tpcds/src/test/java/org/apache/beam/sdk/tpcds/TableSchemaJSONLoaderTest.java +++ b/sdks/java/testing/tpcds/src/test/java/org/apache/beam/sdk/tpcds/TableSchemaJSONLoaderTest.java @@ -42,15 +42,15 @@ public void testStoreReturnsTable() throws Exception { + "sr_reason_sk bigint," + "sr_ticket_number bigint," + "sr_return_quantity bigint," - + "sr_return_amt double," - + "sr_return_tax double," - + "sr_return_amt_inc_tax double," - + "sr_fee double," - + "sr_return_ship_cost double," - + "sr_refunded_cash double," - + "sr_reversed_charge double," - + "sr_store_credit double," - + "sr_net_loss double"; + + "sr_return_amt decimal," + + "sr_return_tax decimal," + + "sr_return_amt_inc_tax decimal," + + "sr_fee decimal," + + "sr_return_ship_cost decimal," + + "sr_refunded_cash decimal," + + "sr_reversed_charge decimal," + + "sr_store_credit decimal," + + "sr_net_loss decimal"; assertEquals(expected, storeReturnsSchemaString); } @@ -63,8 +63,8 @@ public void testItemTable() throws Exception { + "i_rec_start_date varchar," + "i_rec_end_date varchar," + "i_item_desc varchar," - + "i_current_price double," - + "i_wholesale_cost double," + + "i_current_price decimal," + + "i_wholesale_cost decimal," + "i_brand_id bigint," + "i_brand varchar," + "i_class_id bigint," @@ -135,7 +135,7 @@ public void testWarehouseTable() throws Exception { + "w_state varchar," + "w_zip varchar," + "w_country varchar," - + "w_gmt_offset double"; + + "w_gmt_offset decimal"; assertEquals(expected, warehouseSchemaString); } diff --git a/sdks/java/testing/tpcds/src/test/java/org/apache/beam/sdk/tpcds/TpcdsParametersReaderTest.java b/sdks/java/testing/tpcds/src/test/java/org/apache/beam/sdk/tpcds/TpcdsParametersReaderTest.java index 9b9c352541ee..f49e4e1d1d9a 100644 --- a/sdks/java/testing/tpcds/src/test/java/org/apache/beam/sdk/tpcds/TpcdsParametersReaderTest.java +++ b/sdks/java/testing/tpcds/src/test/java/org/apache/beam/sdk/tpcds/TpcdsParametersReaderTest.java @@ -64,6 +64,15 @@ public void testGetAndCheckAllQueries() throws Exception { Assert.assertArrayEquals(expected.toArray(new String[0]), queryNames); } + @Test + public void testGetAndCheckSupportedQueries() throws Exception { + TpcdsOptions tpcdsOptionsSupported = PipelineOptionsFactory.as(TpcdsOptions.class); + tpcdsOptionsSupported.setQueries("supported"); + String[] queryNames = TpcdsParametersReader.getAndCheckQueryNames(tpcdsOptionsSupported); + List expected = TpcdsParametersReader.SUPPORTED_QUERY_NAMES; + Assert.assertArrayEquals(expected.toArray(new String[0]), queryNames); + } + @Test public void testGetAndCheckSpecifiedQueries() throws Exception { String[] queryNames = TpcdsParametersReader.getAndCheckQueryNames(tpcdsOptions); diff --git a/sdks/java/testing/tpcds/src/test/java/org/apache/beam/sdk/tpcds/TpcdsSchemasTest.java b/sdks/java/testing/tpcds/src/test/java/org/apache/beam/sdk/tpcds/TpcdsSchemasTest.java index 9f2397cf945c..fdf9fea95aa2 100644 --- a/sdks/java/testing/tpcds/src/test/java/org/apache/beam/sdk/tpcds/TpcdsSchemasTest.java +++ b/sdks/java/testing/tpcds/src/test/java/org/apache/beam/sdk/tpcds/TpcdsSchemasTest.java @@ -68,8 +68,8 @@ public void testCallCenterSchema() { .addNullableField("cc_state", Schema.FieldType.STRING) .addNullableField("cc_zip", Schema.FieldType.STRING) .addNullableField("cc_country", Schema.FieldType.STRING) - .addNullableField("cc_gmt_offset", Schema.FieldType.DOUBLE) - .addNullableField("cc_tax_percentage", Schema.FieldType.DOUBLE) + .addNullableField("cc_gmt_offset", Schema.FieldType.DECIMAL) + .addNullableField("cc_tax_percentage", Schema.FieldType.DECIMAL) .build(); assertNotEquals(schemaMap.get("call_center"), callCenterSchema); @@ -111,7 +111,7 @@ public void testCustomerAddressSchemaNullable() { .addNullableField("ca_state", Schema.FieldType.STRING) .addNullableField("ca_zip", Schema.FieldType.STRING) .addNullableField("ca_country", Schema.FieldType.STRING) - .addNullableField("ca_gmt_offset", Schema.FieldType.DOUBLE) + .addNullableField("ca_gmt_offset", Schema.FieldType.DECIMAL) .addNullableField("ca_location_type", Schema.FieldType.STRING) .build(); diff --git a/sdks/python/apache_beam/testing/benchmarks/tpcds/README.md b/sdks/python/apache_beam/testing/benchmarks/tpcds/README.md new file mode 100644 index 000000000000..e77f50c3b57c --- /dev/null +++ b/sdks/python/apache_beam/testing/benchmarks/tpcds/README.md @@ -0,0 +1,57 @@ + + +# Apache Beam Python TPC-DS Benchmark & Evaluation Suite + +This package provides a Python TPC-DS benchmark and correctness suite that uses: +- **Cross-language Beam SQL Expansion Service** (`:sdks:java:extensions:sql:expansion-service:shadowJar`) via `apache_beam.transforms.sql.SqlTransform` +- **Prism Runner** via `PrismRunner` +- **Java TPC-DS schemas and queries** copied into `tpcds/build` via Gradle (`copyTpcdsResources`) +- **GCS TPC-DS datasets** (`gs://beam-tpcds/datasets/text/10MB`, `1GB`, `10GB`, `100GB`) + +## Running via Gradle + +The Gradle tasks in `:sdks:python:apache_beam:testing:benchmarks:tpcds` automatically copy the Java TPC-DS schemas and queries into `tpcds/build` and build the Java SQL expansion service shadow jar before running Python. + +### 1. Run the Supported TPC-DS Queries on the 10MB Dataset + +```bash +./gradlew :sdks:python:apache_beam:testing:benchmarks:tpcds:run \ + -Ptpcds.args="--queries=supported --data_size=10MB" +``` + +### 2. Run the 13 Canonical Multi-Table Queries + +```bash +./gradlew :sdks:python:apache_beam:testing:benchmarks:tpcds:run \ + -Ptpcds.args="--queries=canonical13 --data_size=10MB" +``` + +### 3. Run Specific Queries on 1GB GCS Data + +```bash +./gradlew :sdks:python:apache_beam:testing:benchmarks:tpcds:run \ + -Ptpcds.args="--queries=3,7,42,43,52,55,96 --data_size=1GB --results_directory=/tmp/tpcds_results" +``` + +### 4. Run the Pytest Evaluation Suite + +```bash +./gradlew :sdks:python:apache_beam:testing:benchmarks:tpcds:pytest +``` diff --git a/sdks/python/apache_beam/testing/benchmarks/tpcds/__init__.py b/sdks/python/apache_beam/testing/benchmarks/tpcds/__init__.py new file mode 100644 index 000000000000..a0174df802d1 --- /dev/null +++ b/sdks/python/apache_beam/testing/benchmarks/tpcds/__init__.py @@ -0,0 +1,18 @@ +# +# Licensed to the Apache Software Foundation (ASF) under one or more +# contributor license agreements. See the NOTICE file distributed with +# this work for additional information regarding copyright ownership. +# The ASF licenses this file to You under the Apache License, Version 2.0 +# (the "License"); you may not use this file except in compliance with +# the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# + +"""TPC-DS benchmark and evaluation suite for Apache Beam Python SQL.""" diff --git a/sdks/python/apache_beam/testing/benchmarks/tpcds/build.gradle b/sdks/python/apache_beam/testing/benchmarks/tpcds/build.gradle new file mode 100644 index 000000000000..8dac81f94495 --- /dev/null +++ b/sdks/python/apache_beam/testing/benchmarks/tpcds/build.gradle @@ -0,0 +1,84 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * License); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an AS IS BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +apply plugin: org.apache.beam.gradle.BeamModulePlugin +applyPythonNature() + +description = "Apache Beam :: SDKs :: Python :: Benchmarks :: TPC-DS" + +def tpcdsArgsProperty = "tpcds.args" + +def getTpcdsArgs = { + def tpcdsArgsStr = project.findProperty(tpcdsArgsProperty) ?: "" + def tpcdsArgsList = new ArrayList() + if (tpcdsArgsStr.trim()) { + Collections.addAll(tpcdsArgsList, tpcdsArgsStr.trim().split("\\s+")) + } + return tpcdsArgsList +} + +def copyTpcdsResources = tasks.register("copyTpcdsResources", Copy) { + from "${project.rootDir}/sdks/java/testing/tpcds/src/main/resources" + into "${project.buildDir}" +} + +tasks.register("run", Exec) { + dependsOn installGcpTest + dependsOn copyTpcdsResources + dependsOn ":sdks:java:extensions:sql:expansion-service:shadowJar" + + setWorkingDir "${project.rootDir}/sdks/python" + doFirst { + def argsList = getTpcdsArgs() + commandLine 'sh', '-c', ". ${envdir}/bin/activate && python -m apache_beam.testing.benchmarks.tpcds.tpcds_launcher " + String.join(" ", argsList) + } +} + +tasks.register("checkAllQueries", Exec) { + description = "Runs the full 103-query TPC-DS suite on PrismRunner with 10MB dataset and verifies against KNOWN_FAILING_QUERIES." + dependsOn installGcpTest + dependsOn copyTpcdsResources + dependsOn ":sdks:java:extensions:sql:expansion-service:shadowJar" + + setWorkingDir "${project.rootDir}/sdks/python" + doFirst { + def argsList = getTpcdsArgs() + if (!argsList.any { it.startsWith("--queries") }) { + argsList.add("--queries=all") + } + if (!argsList.any { it.startsWith("--check_regression") }) { + argsList.add("--check_regression") + } + commandLine 'sh', '-c', ". ${envdir}/bin/activate && python -m apache_beam.testing.benchmarks.tpcds.tpcds_launcher " + String.join(" ", argsList) + } +} + +tasks.register("pytest", Exec) { + dependsOn installGcpTest + dependsOn copyTpcdsResources + dependsOn ":sdks:java:extensions:sql:expansion-service:shadowJar" + + setWorkingDir "${project.rootDir}/sdks/python" + doFirst { + commandLine 'sh', '-c', ". ${envdir}/bin/activate && python -m pytest -v apache_beam/testing/benchmarks/tpcds/tpcds_test.py" + } +} + +tasks.named("test") { + dependsOn "pytest" +} diff --git a/sdks/python/apache_beam/testing/benchmarks/tpcds/tpcds_launcher.py b/sdks/python/apache_beam/testing/benchmarks/tpcds/tpcds_launcher.py new file mode 100644 index 000000000000..280226cd00e2 --- /dev/null +++ b/sdks/python/apache_beam/testing/benchmarks/tpcds/tpcds_launcher.py @@ -0,0 +1,460 @@ +# +# Licensed to the Apache Software Foundation (ASF) under one or more +# contributor license agreements. See the NOTICE file distributed with +# this work for additional information regarding copyright ownership. +# The ASF licenses this file to You under the Apache License, Version 2.0 +# (the "License"); you may not use this file except in compliance with +# the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# + +"""TPC-DS launcher for Apache Beam Python using cross-language SqlTransform and PrismRunner. + +Example usage via Gradle: + ./gradlew :sdks:python:apache_beam:testing:benchmarks:tpcds:run \\ + -Ptpcds.args="--queries=supported --data_size=10MB" +""" + +import argparse +import contextlib +import dataclasses +import datetime +import glob +import json +import logging +import os +import tempfile +import time +from typing import Any +from typing import Dict +from typing import Iterator +from typing import List +from typing import Optional +from typing import Sequence +from typing import Set + +import apache_beam as beam +from apache_beam.options.pipeline_options import PipelineOptions +from apache_beam.portability.api import schema_pb2 +from apache_beam.testing.benchmarks.tpcds.tpcds_queries import KNOWN_FAILING_QUERIES +from apache_beam.testing.benchmarks.tpcds.tpcds_queries import PLANNER_TIMEOUT_QUERIES +from apache_beam.testing.benchmarks.tpcds.tpcds_queries import get_query_tables +from apache_beam.testing.benchmarks.tpcds.tpcds_queries import parse_query_selection +from apache_beam.testing.benchmarks.tpcds.tpcds_queries import read_query +from apache_beam.testing.benchmarks.tpcds.tpcds_schemas import get_all_table_names +from apache_beam.testing.benchmarks.tpcds.tpcds_schemas import get_table_row_class +from apache_beam.testing.benchmarks.tpcds.tpcds_schemas import parse_csv_line +from apache_beam.transforms.sql import SqlTransform +from apache_beam.typehints import schemas as beam_schemas +from apache_beam.utils import subprocess_server + +_LOGGER = logging.getLogger(__name__) + +DEFAULT_DATA_DIRECTORY = "gs://beam-tpcds/datasets/text" +SUPPORTED_DATA_SIZES: Set[str] = { + "10M", + "10MB", + "1G", + "1GB", + "10G", + "10GB", + "100G", + "100GB", +} + + +def check_and_normalize_data_size(data_size: str) -> str: + """Validates data_size against SUPPORTED_DATA_SIZES and normalizes '10M' -> '10MB', etc.""" + cleaned = data_size.strip().upper() + if cleaned not in SUPPORTED_DATA_SIZES: + raise ValueError( + f"The data size {data_size!r} has not been supported. " + f"Supported sizes: {sorted(SUPPORTED_DATA_SIZES)}") + if cleaned.endswith(("M", "G")): + return f"{cleaned}B" + return cleaned + + +# TODO(https://github.com/apache/beam/issues/40698) remove this contextmanager once Python +# Row supports arbitrary field name +@contextlib.contextmanager +def _sql_schema_field_sanitizer() -> Iterator[None]: + """Temporarily allows NamedTuple schema decoding for SQL aliases like EXPR$n or >120 days.""" + orig_fn = beam_schemas.SchemaTranslation.named_tuple_from_schema + + def _sanitized_named_tuple_from_schema(self, schema): + if any(not f.name.isidentifier() for f in schema.fields): + copied = schema_pb2.Schema() + copied.CopyFrom(schema) + for field in copied.fields: + if not field.name.isidentifier(): + sanitized = "".join( + ch if (ch.isalnum() or ch == "_") else "_" for ch in field.name) + if not sanitized or sanitized[0].isdigit(): + sanitized = f"col_{sanitized}" + field.name = sanitized + schema = copied + return orig_fn(self, schema) + + beam_schemas.SchemaTranslation.named_tuple_from_schema = ( + _sanitized_named_tuple_from_schema) + try: + yield + finally: + beam_schemas.SchemaTranslation.named_tuple_from_schema = orig_fn + + +@dataclasses.dataclass +class TpcdsQueryResult: + query_name: str + job_name: str + data_size: str + dialect: str + is_successful: bool + status: str + start_time: str + end_time: str + elapsed_sec: float + row_count: int = 0 + error_message: str = "" + + +def build_table_pcollections( + pipeline: beam.Pipeline, + table_names: Sequence[str], + data_directory: str = DEFAULT_DATA_DIRECTORY, + data_size: str = "10MB") -> Dict[str, beam.PCollection]: + """Creates schema-typed input PCollections from TPC-DS .dat files.""" + normalized_size = check_and_normalize_data_size(data_size) + tables: Dict[str, beam.PCollection] = {} + base_dir = data_directory.rstrip("/") + for table_name in table_names: + row_cls = get_table_row_class(table_name) + file_pattern = f"{base_dir}/{normalized_size}/{table_name}.dat" + tables[table_name] = ( + pipeline + | f"Read_{table_name}" >> beam.io.ReadFromText(file_pattern) + | f"Parse_{table_name}" >> beam.FlatMap( + lambda line, t=table_name: + [r for r in [parse_csv_line(t, line)] + if r is not None]).with_output_types(row_cls)) + return tables + + +def run_single_query( + query_name: str, + query_sql: str, + pipeline_args: Sequence[str], + data_directory: str = DEFAULT_DATA_DIRECTORY, + data_size: str = "10MB", + results_directory: Optional[str] = None, + skip_planner_timeouts: bool = False) -> TpcdsQueryResult: + """Executes a single TPC-DS query via SqlTransform.""" + normalized_size = check_and_normalize_data_size(data_size) + start_dt = datetime.datetime.now(datetime.timezone.utc) + t0 = time.monotonic() + job_name = f"{query_name}_{int(time.time() * 1000)}" + if skip_planner_timeouts and query_name in PLANNER_TIMEOUT_QUERIES: + err_msg = KNOWN_FAILING_QUERIES[query_name] + end_dt = datetime.datetime.now(datetime.timezone.utc) + return TpcdsQueryResult( + query_name=query_name, + job_name=job_name, + data_size=normalized_size, + dialect="Calcite", + is_successful=False, + status="Failed (PLANNING)", + start_time=start_dt.strftime("%Y-%m-%dT%H:%M:%SZ"), + end_time=end_dt.strftime("%Y-%m-%dT%H:%M:%SZ"), + elapsed_sec=0.0, + error_message=err_msg, + ) + table_names = get_query_tables(query_sql) + + opts_list = list(pipeline_args) + if not any(a.startswith("--runner") for a in opts_list): + opts_list.append("--runner=PrismRunner") + + options = PipelineOptions(opts_list) + + with tempfile.TemporaryDirectory(prefix=f"tpcds_{query_name}_") as tmp_dir: + if results_directory: + out_dir = os.path.join(results_directory, normalized_size) + os.makedirs(out_dir, exist_ok=True) + out_prefix = os.path.join(out_dir, query_name) + else: + out_prefix = os.path.join(tmp_dir, "out") + stage = "PLANNING" + try: + with _sql_schema_field_sanitizer(): + with beam.Pipeline(options=options) as p: + tables = build_table_pcollections( + p, + table_names=table_names, + data_directory=data_directory, + data_size=normalized_size, + ) + sql_input = tables if len(tables) > 1 else next(iter(tables.values())) + sql_out = sql_input | f"Sql_{query_name}" >> SqlTransform(query_sql) + _ = ( + sql_out + | f"Format_{query_name}" >> beam.Map(str) + | f"Write_{query_name}" >> beam.io.WriteToText( + out_prefix, file_name_suffix=".txt", num_shards=1)) + stage = "RUNTIME" + + row_count = 0 + for fpath in sorted(glob.glob(f"{out_prefix}*.txt")): + with open(fpath, "r", encoding="utf-8") as f: + row_count += sum(1 for line in f if line.strip()) + + elapsed = time.monotonic() - t0 + end_dt = datetime.datetime.now(datetime.timezone.utc) + + is_successful = row_count > 0 + if not is_successful: + _LOGGER.warning( + "Output records counter for job %r is %d", job_name, row_count) + + return TpcdsQueryResult( + query_name=query_name, + job_name=job_name, + data_size=normalized_size, + dialect="Calcite", + is_successful=is_successful, + status="Successful" if is_successful else "Failed (EMPTY_RESULT)", + start_time=start_dt.strftime("%Y-%m-%dT%H:%M:%SZ"), + end_time=end_dt.strftime("%Y-%m-%dT%H:%M:%SZ"), + elapsed_sec=round(elapsed, 3), + row_count=row_count, + error_message="" if is_successful else "Returned 0 rows", + ) + except Exception as exc: # pylint: disable=broad-except + elapsed = time.monotonic() - t0 + end_dt = datetime.datetime.now(datetime.timezone.utc) + err_str = str(exc) + first_line = err_str.splitlines()[0] if err_str else repr(exc) + for line in err_str.splitlines(): + if "UnsupportedOperationException" in line or "CannotPlanException" in line: + first_line = line.strip() + break + _LOGGER.error("%s failed at %s: %s", query_name, stage, first_line) + return TpcdsQueryResult( + query_name=query_name, + job_name=job_name, + data_size=normalized_size, + dialect="Calcite", + is_successful=False, + status=f"Failed ({stage})", + start_time=start_dt.strftime("%Y-%m-%dT%H:%M:%SZ"), + end_time=end_dt.strftime("%Y-%m-%dT%H:%M:%SZ"), + elapsed_sec=round(elapsed, 3), + error_message=first_line, + ) + + +def format_summary_table(results: Sequence[TpcdsQueryResult]) -> str: + """Formats the TPC-DS execution summary table.""" + headers = [ + "Query Name", + "Data Size", + "Dialect", + "Status", + "Rows", + "Elapsed(s)", + "Details / Notes", + ] + rows = [] + for r in results: + note = f"Returned {r.row_count} row(s)" if r.is_successful else r.error_message + if len(note) > 85: + note = note[:82] + "..." + rows.append([ + r.query_name, + r.data_size, + r.dialect, + r.status, + str(r.row_count) if r.is_successful else "-", + f"{r.elapsed_sec:.2f}", + note, + ]) + + widths = [len(h) for h in headers] + for row in rows: + for i, cell in enumerate(row): + widths[i] = max(widths[i], len(cell)) + + sep = "+" + "+".join("-" * (w + 2) for w in widths) + "+" + header_line = "|" + "|".join( + f" {h.ljust(widths[i])} " for i, h in enumerate(headers)) + "|" + lines = ["\nTPC-DS Query Execution Summary:", sep, header_line, sep] + for row in rows: + lines.append( + "|" + "|".join(f" {c.ljust(widths[i])} " + for i, c in enumerate(row)) + "|") + lines.append(sep) + return "\n".join(lines) + + +def main(argv: Optional[Sequence[str]] = None) -> List[TpcdsQueryResult]: + parser = argparse.ArgumentParser( + description="Run Apache Beam Python TPC-DS benchmark suite.") + parser.add_argument( + "--queries", + default="supported", + help= + "Comma-separated query numbers (e.g. '3,7,96'), 'canonical13', 'supported', or 'all'.", + ) + parser.add_argument( + "--data_directory", + default=DEFAULT_DATA_DIRECTORY, + help= + "Path or GCS URI containing TPC-DS text data files (/.dat).", + ) + parser.add_argument( + "--data_size", + default="10MB", + help= + "Dataset scale factor ('10M', '10MB', '1G', '1GB', '10G', '10GB', '100G', '100GB').", + ) + parser.add_argument( + "--queries_directory", + default=None, + help="Optional path to directory containing query*.sql files.", + ) + parser.add_argument( + "--results_directory", + default=None, + help="Optional directory to write query results.", + ) + parser.add_argument( + "--summary_json", + default=None, + help="Optional file path to write JSON execution summary.", + ) + parser.add_argument( + "--check_regression", + default=False, + action=argparse.BooleanOptionalAction, + help="Verify query outcomes against KNOWN_FAILING_QUERIES. Queries not in " + "KNOWN_FAILING_QUERIES must succeed; queries in KNOWN_FAILING_QUERIES must " + "fail with an error matching their expected error substring.", + ) + + known_args, pipeline_args = parser.parse_known_args(argv) + # Fail early if tpcds/build/schemas has not been populated by Gradle. + get_all_table_names() + normalized_size = check_and_normalize_data_size(known_args.data_size) + query_names = parse_query_selection(known_args.queries) + + all_results: List[TpcdsQueryResult] = [] + with subprocess_server.SubprocessServer.cache_subprocesses(): + for qname in query_names: + query_sql = read_query(qname, known_args.queries_directory) + res = run_single_query( + query_name=qname, + query_sql=query_sql, + pipeline_args=pipeline_args, + data_directory=known_args.data_directory, + data_size=normalized_size, + results_directory=known_args.results_directory, + skip_planner_timeouts=known_args.check_regression, + ) + all_results.append(res) + + print(format_summary_table(all_results)) + + if known_args.summary_json: + os.makedirs( + os.path.dirname(os.path.abspath(known_args.summary_json)), + exist_ok=True) + with open(known_args.summary_json, "w", encoding="utf-8") as f: + json.dump( + [dataclasses.asdict(r) for r in all_results], + f, + indent=2, + default=str, + ) + + if known_args.check_regression: + verify_no_regressions(all_results) + + return all_results + + +def verify_no_regressions( + results: Sequence[TpcdsQueryResult], + known_failing: Optional[Dict[str, str]] = None) -> None: + """Verifies that query results match SUPPORTED_QUERIES and KNOWN_FAILING_QUERIES. + + Raises RuntimeError if: + - A query in SUPPORTED_QUERIES fails, or + - A query in KNOWN_FAILING_QUERIES fails with an unexpected error message, or + - A query is unclassified (in neither SUPPORTED_QUERIES nor KNOWN_FAILING_QUERIES). + + Queries in KNOWN_FAILING_QUERIES that succeed are tolerated with a warning. + """ + expected_map = KNOWN_FAILING_QUERIES if known_failing is None else known_failing + supported_set = set(parse_query_selection("supported")) + regressions: List[str] = [] + unexpected_passes: List[str] = [] + expected_failures_matched = 0 + passed_queries = 0 + + for r in results: + expected_err = expected_map.get(r.query_name) + if expected_err is not None: + if r.is_successful: + warn_msg = ( + f"[UNEXPECTED PASS] {r.query_name} is listed in " + f"KNOWN_FAILING_QUERIES ({expected_err!r}) but succeeded! " + f"Consider moving it to SUPPORTED_QUERIES if non-flaky.") + _LOGGER.warning(warn_msg) + unexpected_passes.append(warn_msg) + elif expected_err not in r.error_message: + regressions.append( + f"[ERROR MISMATCH] {r.query_name} failed with " + f"unexpected error: {r.error_message!r} (expected substring: {expected_err!r})" + ) + else: + expected_failures_matched += 1 + elif r.query_name in supported_set: + if not r.is_successful: + regressions.append( + f"[REGRESSION] {r.query_name} is in SUPPORTED_QUERIES " + f"but failed with: {r.error_message}") + else: + passed_queries += 1 + else: + regressions.append( + f"[UNCLASSIFIED] {r.query_name} is in neither SUPPORTED_QUERIES " + f"nor KNOWN_FAILING_QUERIES (status={r.status}, error={r.error_message!r})" + ) + + _LOGGER.info( + "TPC-DS Regression Check Summary: %d passed as expected, " + "%d known failures matched expected error, " + "%d unexpected pass warning(s), " + "%d regression(s).", + passed_queries, + expected_failures_matched, + len(unexpected_passes), + len(regressions)) + if regressions: + for msg in regressions: + _LOGGER.error(" - %s", msg) + raise RuntimeError( + f"TPC-DS regression check failed with {len(regressions)} issue(s):\n" + + "\n".join(regressions)) + + +if __name__ == "__main__": + logging.basicConfig(level=logging.INFO) + main() diff --git a/sdks/python/apache_beam/testing/benchmarks/tpcds/tpcds_queries.py b/sdks/python/apache_beam/testing/benchmarks/tpcds/tpcds_queries.py new file mode 100644 index 000000000000..e7d45fc57d48 --- /dev/null +++ b/sdks/python/apache_beam/testing/benchmarks/tpcds/tpcds_queries.py @@ -0,0 +1,195 @@ +# +# Licensed to the Apache Software Foundation (ASF) under one or more +# contributor license agreements. See the NOTICE file distributed with +# this work for additional information regarding copyright ownership. +# The ASF licenses this file to You under the Apache License, Version 2.0 +# (the "License"); you may not use this file except in compliance with +# the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# + +"""TPC-DS SQL query loader and table identifier extraction.""" + +import os +import re +from typing import Dict +from typing import List +from typing import Optional +from typing import Set + +from apache_beam.testing.benchmarks.tpcds.tpcds_schemas import TPCDS_BUILD_DIR +from apache_beam.testing.benchmarks.tpcds.tpcds_schemas import get_all_table_names + +TPCDS_QUERIES_DIR = os.path.join(TPCDS_BUILD_DIR, "queries") + +# The 13 canonical multi-table TPC-DS star/snowflake aggregate queries. +CANONICAL_13_QUERIES: List[str] = [ + "3", "7", "13", "15", "19", "26", "42", "43", "48", "52", "55", "91", "96" +] + +# The 55 TPC-DS queries supported by Beam Calcite SQL (returning > 0 rows). +# Together with KNOWN_FAILING_QUERIES (48 queries), this partitions all 103 TPC-DS queries. +# yapf: disable +SUPPORTED_QUERIES: List[str] = [ + "1", "3", "4", "7", "9", "10", "11", "12", "15", "17", + "19", "20", "21", "26", "31", "32", "33", "34", "35", "38", + "40", "42", "43", "45", "46", "49", "50", "52", "53", "55", + "56", "59", "60", "61", "63", "65", "66", "68", "69", "71", + "73", "74", "75", "76", "79", "82", "84", "87", "88", "89", + "92", "93", "96", "97", "99", +] +# yapf: enable + +# The 48 known failing TPC-DS queries on Beam Calcite SQL with their expected error substring. +# Used by `--check_regression` and `:sdks:python:apache_beam:testing:benchmarks:tpcds:checkAllQueries` +# to verify that all 55 SUPPORTED_QUERIES succeed on PrismRunner and that the remaining 48 known +# failing queries fail with their expected diagnostic without regressing. +# pylint: disable=line-too-long +KNOWN_FAILING_QUERIES: Dict[str, str] = { + "query2": "UnsupportedOperationException: Only support column reference or struct field access in conjunction clause", + "query5": "Error while applying rule BeamBasicAggregationRule", + "query6": "VolcanoPlanner timeout (correlated subquery + SINGLE_VALUE join permutation search)", + "query8": "UnsupportedOperationException: Only support column reference or struct field access in conjunction clause", + "query13": "UnsupportedOperationException: Operator OR is not supported in join condition", + "query14a": "Error while applying rule BeamBasicAggregationRule", + "query14b": "UnsupportedOperationException: Aggregator [SINGLE_VALUE] is not supported", + "query16": "UnsupportedOperationException: Does not support COUNT DISTINCT", + "query18": "Error while applying rule BeamBasicAggregationRule", + "query22": "Error while applying rule BeamBasicAggregationRule", + "query23a": "UnsupportedOperationException: Aggregator [SINGLE_VALUE] is not supported", + "query23b": "UnsupportedOperationException: Aggregator [SINGLE_VALUE] is not supported", + "query24a": "VolcanoPlanner timeout (>12-table join permutation search)", + "query24b": "VolcanoPlanner timeout (>12-table join permutation search)", + "query25": "Returned 0 rows", + "query27": "Error while applying rule BeamBasicAggregationRule", + "query28": "ParseException: Unable to parse query", + "query29": "Returned 0 rows", + "query30": "UnsupportedOperationException: Cross join with a residual (non-equi) condition is not supported yet.", + "query36": "Error while applying rule BeamBasicAggregationRule", + "query37": "Returned 0 rows", + "query39a": "Returned 0 rows", + "query39b": "Returned 0 rows", + "query41": "ParseException: Unable to parse query", + "query44": "UnsupportedOperationException: Aggregator [SINGLE_VALUE] is not supported", + "query47": "IllegalArgumentException: DECIMAL NOT NULL is not nullable", + "query48": "UnsupportedOperationException: Operator OR is not supported in join condition", + "query51": "IllegalArgumentException: DECIMAL NOT NULL is not nullable", + "query54": "UnsupportedOperationException: Aggregator [SINGLE_VALUE] is not supported", + "query57": "Returned 0 rows", + "query58": "UnsupportedOperationException: Aggregator [SINGLE_VALUE] is not supported", + "query62": "Error while applying rule JoinPushExpressionsRule", + "query64": "VolcanoPlanner timeout (36-table join permutation search)", + "query67": "Error while applying rule BeamBasicAggregationRule", + "query70": "Error while applying rule BeamBasicAggregationRule", + "query72": "ParseException: Unable to parse query", + "query77": "Error while applying rule BeamBasicAggregationRule", + "query78": "CalcFn failed to evaluate", + "query80": "Error while applying rule BeamBasicAggregationRule", + "query81": "UnsupportedOperationException: Cross join with a residual (non-equi) condition is not supported yet.", + "query83": "Error decoding input stream with coder", + "query85": "Returned 0 rows", + "query86": "Error while applying rule BeamBasicAggregationRule", + "query90": "CalcFn failed to evaluate", + "query91": "VolcanoPlanner timeout (7-table join permutation search)", + "query94": "Error while applying rule JoinPushExpressionsRule", + "query95": "Error while applying rule JoinPushExpressionsRule", + "query98": "IllegalArgumentException: DECIMAL NOT NULL is not nullable", +} +# pylint: enable=line-too-long + +# Subset of KNOWN_FAILING_QUERIES where Calcite's VolcanoPlanner explores exponential join +# permutations (7 to 36 table joins) before failing with CannotPlanException. +PLANNER_TIMEOUT_QUERIES: Set[str] = { + "query6", "query24a", "query24b", "query64", "query91" +} + + +def _require_queries_dir(queries_dir: Optional[str] = None) -> str: + target_dir = queries_dir or TPCDS_QUERIES_DIR + if not os.path.isdir(target_dir): + raise RuntimeError( + f"TPC-DS queries directory not found at {target_dir}. " + "Running TPC-DS as a standalone apache-beam module is not supported; " + "run via Gradle (:sdks:python:apache_beam:testing:benchmarks:tpcds:copyTpcdsResources) " + "to populate tpcds/build.") + return target_dir + + +def get_all_query_names() -> List[str]: + """Returns all 103 TPC-DS query names (query1 .. query99 including a/b variants).""" + names: List[str] = [] + split_queries = {14, 23, 24, 39} + for i in range(1, 100): + if i in split_queries: + names.append(f"query{i}a") + names.append(f"query{i}b") + else: + names.append(f"query{i}") + return names + + +def normalize_query_name(q: str) -> str: + """Normalizes '3', '03', 'q03', 'query3' -> 'query3'.""" + cleaned = q.strip().lower() + m = re.match(r"^(?:query|q)?0*(\d+[ab]?)$", cleaned) + if m: + return f"query{m.group(1)}" + return cleaned + + +def parse_query_selection(queries_arg: str) -> List[str]: + """Parses a --queries argument ('canonical13', 'supported', 'all', or comma-separated list).""" + cleaned = queries_arg.strip().lower() + if cleaned in ("canonical13", "canonical"): + return [normalize_query_name(q) for q in CANONICAL_13_QUERIES] + if cleaned == "supported": + return [normalize_query_name(q) for q in SUPPORTED_QUERIES] + if cleaned == "all": + return get_all_query_names() + all_valid = set(get_all_query_names()) + result = [] + for part in re.split(r"[\s,]+", queries_arg.strip()): + if not part: + continue + qname = normalize_query_name(part) + if qname not in all_valid: + raise ValueError(f"Unsupported TPC-DS query name: {part!r} ({qname})") + result.append(qname) + return result + + +def read_query(query_name: str, queries_dir: Optional[str] = None) -> str: + """Reads a TPC-DS SQL query by name (e.g. 'query3' or '3') from tpcds/build/queries.""" + qname = normalize_query_name(query_name) + target_dir = _require_queries_dir(queries_dir) + candidate = os.path.join(target_dir, f"{qname}.sql") + if not os.path.isfile(candidate): + raise RuntimeError( + f"Could not find SQL file for {qname} at {candidate}. " + "Run via Gradle (:sdks:python:apache_beam:testing:benchmarks:tpcds:copyTpcdsResources)." + ) + with open(candidate, "r", encoding="utf-8") as f: + return _strip_sql_comments(f.read()) + + +def _strip_sql_comments(sql: str) -> str: + lines = [] + for line in sql.splitlines(): + if line.lstrip().startswith("--"): + continue + lines.append(line) + return "\n".join(lines).strip() + + +def get_query_tables(query_sql: str) -> List[str]: + """Extracts the TPC-DS table names referenced in the given SQL query.""" + tokens: Set[str] = set( + re.findall(r"\b[a-zA-Z_][a-zA-Z0-9_]*\b", query_sql.lower())) + return [t for t in get_all_table_names() if t.lower() in tokens] diff --git a/sdks/python/apache_beam/testing/benchmarks/tpcds/tpcds_schemas.py b/sdks/python/apache_beam/testing/benchmarks/tpcds/tpcds_schemas.py new file mode 100644 index 000000000000..40831972a8e1 --- /dev/null +++ b/sdks/python/apache_beam/testing/benchmarks/tpcds/tpcds_schemas.py @@ -0,0 +1,141 @@ +# +# Licensed to the Apache Software Foundation (ASF) under one or more +# contributor license agreements. See the NOTICE file distributed with +# this work for additional information regarding copyright ownership. +# The ASF licenses this file to You under the Apache License, Version 2.0 +# (the "License"); you may not use this file except in compliance with +# the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# + +"""TPC-DS table schemas and row converters for Beam Python SQL.""" + +import decimal +import json +import os +import typing +from typing import Any +from typing import Dict +from typing import List +from typing import Optional +from typing import Tuple +from typing import Type + +from apache_beam import coders + +TPCDS_BUILD_DIR = os.path.join( + os.path.dirname(os.path.abspath(__file__)), "build") +TPCDS_SCHEMAS_DIR = os.path.join(TPCDS_BUILD_DIR, "schemas") + +_SCHEMA_CACHE: Dict[str, List[Tuple[str, str]]] = {} +_ROW_CLASS_CACHE: Dict[str, Type[Tuple[Any, ...]]] = {} + + +def _require_schemas_dir() -> str: + if not os.path.isdir(TPCDS_SCHEMAS_DIR): + raise RuntimeError( + f"TPC-DS schema directory not found at {TPCDS_SCHEMAS_DIR}. " + "Running TPC-DS as a standalone apache-beam module is not supported; " + "run via Gradle (:sdks:python:apache_beam:testing:benchmarks:tpcds:copyTpcdsResources) " + "to populate tpcds/build.") + return TPCDS_SCHEMAS_DIR + + +def get_raw_table_schema(table_name: str) -> List[Tuple[str, str]]: + """Returns the list of (column_name, raw_type) for a TPC-DS table from tpcds/build/schemas.""" + if table_name in _SCHEMA_CACHE: + return list(_SCHEMA_CACHE[table_name]) + + schemas_dir = _require_schemas_dir() + json_path = os.path.join(schemas_dir, f"{table_name}.json") + if not os.path.isfile(json_path): + raise RuntimeError( + f"TPC-DS schema file for table {table_name!r} not found at {json_path}. " + "Run via Gradle (:sdks:python:apache_beam:testing:benchmarks:tpcds:copyTpcdsResources)." + ) + with open(json_path, "r", encoding="utf-8") as f: + data = json.load(f) + cols = [(col["name"], col["type"].lower()) for col in data["schema"]] + _SCHEMA_CACHE[table_name] = cols + return list(cols) + + +def get_all_table_names() -> List[str]: + """Returns sorted list of all TPC-DS table names from tpcds/build/schemas.""" + schemas_dir = _require_schemas_dir() + names = sorted( + os.path.splitext(fname)[0] for fname in os.listdir(schemas_dir) + if fname.endswith(".json")) + if not names: + raise RuntimeError(f"No TPC-DS schema JSON files found in {schemas_dir}.") + return names + + +def _python_field_type(raw_type: str) -> Any: + """Maps a raw TPC-DS JSON schema type to a nullable Python type hint for RowCoder.""" + t = raw_type.lower() + if t in ("identifier", "integer"): + return Optional[int] + elif "decimal" in t: + return Optional[decimal.Decimal] + else: + return Optional[str] + + +def get_table_row_class(table_name: str) -> Type[Tuple[Any, ...]]: + """Returns a NamedTuple class registered with RowCoder for the given table.""" + if table_name in _ROW_CLASS_CACHE: + return _ROW_CLASS_CACHE[table_name] + + raw_cols = get_raw_table_schema(table_name) + fields = [(col_name, _python_field_type(col_type)) + for col_name, col_type in raw_cols] + type_name = f"Tpcds_{table_name}" + row_cls = typing.NamedTuple(type_name, fields) + coders.registry.register_coder(row_cls, coders.RowCoder) + _ROW_CLASS_CACHE[table_name] = row_cls + return row_cls + + +def coerce_value(val: Any, raw_type: str) -> Any: + """Coerces a Python or string value to match the target column type.""" + if val is None or val == "": + return None + t = raw_type.lower() + if t in ("identifier", "integer"): + return int(val) + elif "decimal" in t: + if isinstance(val, decimal.Decimal): + return val + if isinstance(val, float): + return decimal.Decimal(f"{val:.2f}") + return decimal.Decimal(str(val)) + else: + return str(val) + + +def parse_csv_line(table_name: str, + line: str, + delimiter: str = "|") -> Optional[Any]: + """Parses a TPC-DS .dat pipe-delimited line into a typed NamedTuple row.""" + stripped = line.rstrip("\r\n") + if not stripped: + return None + raw_cols = get_raw_table_schema(table_name) + row_cls = get_table_row_class(table_name) + parts = stripped.split(delimiter) + # TPC-DS .dat files have a trailing '|' delimiter + if len(parts) == len(raw_cols) + 1 and parts[-1] == "": + parts = parts[:-1] + args = [] + for idx, (_, col_type) in enumerate(raw_cols): + raw_val = parts[idx] if idx < len(parts) else "" + args.append(coerce_value(raw_val, col_type)) + return row_cls(*args) diff --git a/sdks/python/apache_beam/testing/benchmarks/tpcds/tpcds_test.py b/sdks/python/apache_beam/testing/benchmarks/tpcds/tpcds_test.py new file mode 100644 index 000000000000..b19da0f62681 --- /dev/null +++ b/sdks/python/apache_beam/testing/benchmarks/tpcds/tpcds_test.py @@ -0,0 +1,78 @@ +# +# Licensed to the Apache Software Foundation (ASF) under one or more +# contributor license agreements. See the NOTICE file distributed with +# this work for additional information regarding copyright ownership. +# The ASF licenses this file to You under the Apache License, Version 2.0 +# (the "License"); you may not use this file except in compliance with +# the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# + +"""Tests for the Apache Beam Python TPC-DS suite with cross-language SQL and PrismRunner.""" + +import logging +import unittest + +from apache_beam.testing.benchmarks.tpcds import tpcds_launcher +from apache_beam.testing.benchmarks.tpcds import tpcds_queries +from apache_beam.testing.benchmarks.tpcds import tpcds_schemas + + +class TpcdsUnitTest(unittest.TestCase): + """Unit tests for TPC-DS benchmark suite.""" + def test_schema_coercion(self): + self.assertEqual(tpcds_schemas.coerce_value("42", "integer"), 42) + self.assertEqual(tpcds_schemas.coerce_value("123", "identifier"), 123) + self.assertEqual( + str(tpcds_schemas.coerce_value("19.99", "decimal(7,2)")), "19.99") + self.assertEqual(tpcds_schemas.coerce_value("abc", "varchar(20)"), "abc") + self.assertIsNone(tpcds_schemas.coerce_value("", "integer")) + self.assertIsNone(tpcds_schemas.coerce_value(None, "decimal(7,2)")) + + def test_known_failing_queries_map_valid(self): + all_queries = set(tpcds_queries.get_all_query_names()) + supported_queries = set(tpcds_queries.parse_query_selection("supported")) + failing_queries = set(tpcds_queries.KNOWN_FAILING_QUERIES.keys()) + for qname, expected_err in tpcds_queries.KNOWN_FAILING_QUERIES.items(): + self.assertIn(qname, all_queries) + self.assertTrue(len(expected_err) > 0) + self.assertTrue( + tpcds_queries.PLANNER_TIMEOUT_QUERIES.issubset(failing_queries)) + self.assertEqual(supported_queries & failing_queries, set()) + self.assertEqual(supported_queries | failing_queries, all_queries) + # 55 supported + 48 known failing = 103 total queries + self.assertEqual(len(supported_queries), 55) + self.assertEqual(len(failing_queries), 48) + + def test_query_selection_and_normalization(self): + qnames = tpcds_queries.parse_query_selection("canonical13") + self.assertEqual(len(qnames), 13) + self.assertEqual(tpcds_queries.normalize_query_name("03"), "query3") + self.assertEqual(tpcds_queries.normalize_query_name("q14a"), "query14a") + self.assertEqual( + tpcds_queries.parse_query_selection("3, 7, query96"), + ["query3", "query7", "query96"]) + with self.assertRaises(ValueError): + tpcds_queries.parse_query_selection("query999") + + def test_supported_data_sizes(self): + self.assertEqual( + tpcds_launcher.check_and_normalize_data_size("10M"), "10MB") + self.assertEqual( + tpcds_launcher.check_and_normalize_data_size("10MB"), "10MB") + self.assertEqual(tpcds_launcher.check_and_normalize_data_size("1G"), "1GB") + self.assertEqual(tpcds_launcher.check_and_normalize_data_size("1GB"), "1GB") + with self.assertRaises(ValueError): + tpcds_launcher.check_and_normalize_data_size("5G") + + +if __name__ == "__main__": + logging.getLogger().setLevel(logging.INFO) + unittest.main() diff --git a/settings.gradle.kts b/settings.gradle.kts index e1ed639b0d24..1d061e145ff0 100644 --- a/settings.gradle.kts +++ b/settings.gradle.kts @@ -312,6 +312,7 @@ include(":sdks:java:transform-service:controller-container") include(":sdks:python") include(":sdks:python:apache_beam:testing:load_tests") include(":sdks:python:apache_beam:testing:benchmarks:nexmark") +include(":sdks:python:apache_beam:testing:benchmarks:tpcds") include(":sdks:python:container") include(":sdks:python:container:py310") include(":sdks:python:container:py311")