You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@beam.apache.org by yh...@apache.org on 2026/10/22 16:05:07 UTC
(beam) branch master updated: Expand Beam SQL TPC-DS query support and add Python TPC-DS suite (#40695)
This is an automated email from the ASF dual-hosted git repository.
Abacn pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/beam.git
The following commit(s) were added to refs/heads/master by this push:
new 3d5375ac59d Expand Beam SQL TPC-DS query support and add Python TPC-DS suite (#40695)
3d5375ac59d is described below
commit 3d5375ac59d0a93f5537375080eb1dd43637c1d8
Author: Yi Hu <ya...@google.com>
AuthorDate: Thu Oct 22 12:04:59 2026 -0400
Expand Beam SQL TPC-DS query support and add Python TPC-DS suite (#40695)
* Expand Beam SQL TPC-DS query support and add Python TPC-DS suite
- Add `BeamCrossJoinRel` and `BeamCrossJoinRule` to support bounded-by-bounded pure cross joins in Beam Calcite SQL.
- Extend `BeamCoGBKJoinRule` to split mixed equi/non-equi inner join predicates into an equi-join followed by a post-join `BeamCalcRel` filter.
- Support `LAG`, `LEAD`, `FIRST_VALUE`, `LAST_VALUE`, `NTH_VALUE` (`RESPECT`/`IGNORE NULLS`), and `CUME_DIST` analytic window functions in `BeamWindowRel` and `BeamBuiltinAnalyticFunctions`, and convert window inputs to `BeamLogicalConvention` in `BeamWindowRule`.
- Fix `NULLS FIRST` and `NULLS LAST` comparator ordering in `BeamSortRel` when combined with `ASC`/`DESC`.
- Normalize set operator (`UNION`, `INTERSECT`, `EXCEPT`) input schemas in `BeamSetOperatorRelBase` when inputs differ only by field names or nullability.
- Return `NULL` instead of `0.0` from `VarianceFn` for empty or single-element sample variance/stddev groups to match SQL standard semantics.
- Remove workarounds in TPC-DS table schemas used to cast `DECIMAL` to `DOUBLE` for fractional columns
Having these improvements, 55 out of 103 queries now passing (up from 24). Then for test suite,
- add `--queries=supported` and `10M`/`10MB` dataset size support to `TpcdsParametersReader`.
- Add the Python TPC-DS benchmark and regression test suite (`sdks/python/apache_beam/testing/benchmarks/tpcds`) running on `PrismRunner` and cross-language `SqlTransform`.
- Update Java TPC-DS GitHub Actions workflows to use `--queries=supported` and add the `beam_PostCommit_Python_Tpcds_Prism.yml` workflow.
- Move query result file to a temp bucket with lifecycle policy, prevent resource leak
---
.../beam_PostCommit_Java_Tpcds_Dataflow.json | 1 +
.../beam_PostCommit_Java_Tpcds_Flink.json | 2 +-
.../beam_PostCommit_Java_Tpcds_Spark.json | 2 +-
.github/workflows/README.md | 1 +
.../beam_PostCommit_Java_Tpcds_Dataflow.yml | 8 +-
.../workflows/beam_PostCommit_Java_Tpcds_Flink.yml | 4 +-
.../workflows/beam_PostCommit_Java_Tpcds_Spark.yml | 4 +-
....yml => beam_PostCommit_Python_Tpcds_Prism.yml} | 46 +--
.../extensions/sql/impl/CalciteQueryPlanner.java | 5 +
.../extensions/sql/impl/planner/BeamRuleSets.java | 3 +
.../sql/impl/planner/RelMdNodeStats.java | 6 +
.../extensions/sql/impl/rel/BeamCrossJoinRel.java | 173 ++++++++
.../sql/impl/rel/BeamSetOperatorRelBase.java | 94 ++++-
.../sdk/extensions/sql/impl/rel/BeamSortRel.java | 11 +-
.../sdk/extensions/sql/impl/rel/BeamWindowRel.java | 104 ++++-
.../sql/impl/rule/BeamBasicAggregationRule.java | 4 +-
.../sql/impl/rule/BeamCoGBKJoinRule.java | 119 +++++-
...amCoGBKJoinRule.java => BeamCrossJoinRule.java} | 39 +-
.../extensions/sql/impl/rule/BeamWindowRule.java | 4 +-
.../extensions/sql/impl/schema/BeamTableUtils.java | 2 +
.../impl/transform/BeamBuiltinAggregations.java | 10 +-
.../transform/BeamBuiltinAnalyticFunctions.java | 246 ++++++++++-
.../sql/impl/transform/agg/VarianceFn.java | 22 +-
.../extensions/sql/BeamAnalyticFunctionsTest.java | 112 +++++
.../sdk/extensions/sql/BeamSqlDslJoinTest.java | 13 +-
.../rel/BeamCoGBKJoinRelBoundedVsBoundedTest.java | 42 +-
.../org/apache/beam/sdk/tpcds/QueryReader.java | 4 +-
.../beam/sdk/tpcds/TableSchemaJSONLoader.java | 4 +-
.../org/apache/beam/sdk/tpcds/TpcdsOptions.java | 5 +-
.../beam/sdk/tpcds/TpcdsParametersReader.java | 20 +-
.../org/apache/beam/sdk/tpcds/TpcdsSchemas.java | 166 ++++----
.../beam/sdk/tpcds/TableSchemaJSONLoaderTest.java | 24 +-
.../beam/sdk/tpcds/TpcdsParametersReaderTest.java | 9 +
.../apache/beam/sdk/tpcds/TpcdsSchemasTest.java | 6 +-
.../apache_beam/testing/benchmarks/tpcds/README.md | 57 +++
.../testing/benchmarks/tpcds/__init__.py | 18 +
.../testing/benchmarks/tpcds/build.gradle | 84 ++++
.../testing/benchmarks/tpcds/tpcds_launcher.py | 460 +++++++++++++++++++++
.../testing/benchmarks/tpcds/tpcds_queries.py | 195 +++++++++
.../testing/benchmarks/tpcds/tpcds_schemas.py | 141 +++++++
.../testing/benchmarks/tpcds/tpcds_test.py | 78 ++++
settings.gradle.kts | 1 +
42 files changed, 2114 insertions(+), 235 deletions(-)
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 00000000000..0967ef424bc
--- /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 e3d6056a5de..b2683333323 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 e3d6056a5de..b2683333323 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 c41191e4fe3..178fb650268 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`| [](https://github.com/apache/beam/actions/workflows/beam_PostCommit_Python_Examples_Spark.yml?query=event%3Asch [...]
| [ 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`| [](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`| [](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`| [](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`| [](https://github.com/apache/beam/actio [...]
| [ 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`| [ | ['3.9','3.11','3.12'] |`beam_PostCommit_Python_ValidatesRunner_Dataflow.json`| [](https://github.com/apache/beam/actions/workflows/beam_Post [...]
diff --git a/.github/workflows/beam_PostCommit_Java_Tpcds_Dataflow.yml b/.github/workflows/beam_PostCommit_Java_Tpcds_Dataflow.yml
index 35bd18a7d54..e5bc892bb2e 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 f84dea31c34..bb4c3066ddf 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 6ce87271a4b..8f67af86636 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_Java_Tpcds_Flink.yml b/.github/workflows/beam_PostCommit_Python_Tpcds_Prism.yml
similarity index 60%
copy from .github/workflows/beam_PostCommit_Java_Tpcds_Flink.yml
copy to .github/workflows/beam_PostCommit_Python_Tpcds_Prism.yml
index f84dea31c34..168d60b65f3 100644
--- a/.github/workflows/beam_PostCommit_Java_Tpcds_Flink.yml
+++ b/.github/workflows/beam_PostCommit_Python_Tpcds_Prism.yml
@@ -13,13 +13,13 @@
# See the License for the specific language governing permissions and
# limitations under the License.
-name: PostCommit Java Tpcds Flink
+name: PostCommit Python Tpcds Prism
on:
schedule:
- cron: '30 4/6 * * *'
pull_request:
- paths: ['release/trigger_all_tests.json', '.github/trigger_files/beam_PostCommit_Java_Tpcds_Flink.json']
+ paths: ['release/trigger_all_tests.json', '.github/trigger_files/beam_PostCommit_Python_Tpcds_Prism.json']
workflow_dispatch:
# Setting explicit permissions for the action
@@ -45,43 +45,19 @@ concurrency:
env:
DEVELOCITY_ACCESS_KEY: ${{ secrets.BEAM_DEVELOCITY_ACCESS_KEY }}
- GRADLE_COMMAND_ARGUMENTS: |
- --runner=FlinkRunner
- --parallelism=4
- --dataSize=1GB
- --sourceType=PARQUET
- --dataDirectory=gs://beam-tpcds/datasets/parquet/nonpartitioned
- --resultsDirectory=gs://beam-tpcds/results/flink/
- --tpcParallel=1
- INFLUXDB_USER: ${{ secrets.INFLUXDB_USER }}
- INFLUXDB_USER_PASSWORD: ${{ secrets.INFLUXDB_USER_PASSWORD }}
- tpcdsBigQueryArgs: |
- --bigQueryTable=tpcds
- --bigQueryDataset=tpcds
- --project=apache-beam-testing
- --resourceNameMode=QUERY_RUNNER_AND_MODE
- --exportSummaryToBigQuery=true
- --tempLocation=gs://temp-storage-for-perf-tests/tpcds
- tpcdsInfluxDBArgs: |
- --influxDatabase=beam_test_metrics
- --influxHost=http://10.128.0.96:8086
- --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
-
+
jobs:
- beam_PostCommit_Java_Tpcds_Flink:
+ 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: 240
+ timeout-minutes: 60
name: ${{ matrix.job_name }}
strategy:
matrix:
- job_name: ["beam_PostCommit_Java_Tpcds_Flink"]
+ job_name: ["beam_PostCommit_Python_Tpcds_Prism"]
steps:
- uses: actions/checkout@v7
with:
@@ -90,10 +66,10 @@ jobs:
uses: ./.github/actions/setup-action
- name: Setup environment
uses: ./.github/actions/setup-environment-action
- - name: run PostCommit Java Tpcds Flink script
+ 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:java:testing:tpcds:run
- arguments: |
- -Ptpcds.runner=:runners:flink:2.2 \
- "-Ptpcds.args=${{env.tpcdsBigQueryArgs}} ${{env.tpcdsInfluxDBArgs}} ${{ env.GRADLE_COMMAND_ARGUMENTS }} --queries=${{env.tpcdsQueriesArg}}" \
+ 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 cedfc5dc6d1..729ab8f8d84 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 class CalciteQueryPlanner implements QueryPlanner {
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 8382fcb6e38..2e3ef182386 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.BeamBasicAggregationRule;
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 40db8b074ef..51b197a7eac 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 java.util.stream.Collectors;
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 class RelMdNodeStats implements MetadataHandler<NodeStatsMetadata> {
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 00000000000..ebdc7379688
--- /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.
+ *
+ * <p>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).
+ *
+ * <p>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<CorrelationId> 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<PCollectionList<Row>, PCollection<Row>> 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<PCollectionList<Row>, PCollection<Row>> {
+ @Override
+ public PCollection<Row> expand(PCollectionList<Row> pinput) {
+ PCollection<Row> leftRows =
+ pinput
+ .get(0)
+ .apply(
+ "left_TimestampCombiner",
+ Window.<Row>configure().withTimestampCombiner(TimestampCombiner.EARLIEST));
+ PCollection<Row> rightRows =
+ pinput
+ .get(1)
+ .apply(
+ "right_TimestampCombiner",
+ Window.<Row>configure().withTimestampCombiner(TimestampCombiner.EARLIEST));
+
+ final Schema outputSchema = CalciteUtils.toSchema(getRowType());
+ final PCollectionView<List<Row>> 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<Row, Row> {
+ private final Schema outputSchema;
+ private final PCollectionView<List<Row>> rightView;
+ private final boolean leftJoin;
+ private final int rightFieldCount;
+
+ CrossJoinDoFn(
+ Schema outputSchema,
+ PCollectionView<List<Row>> 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<Row> 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 cad845c2b3a..efbdfc4111b 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.PCollection;
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 class BeamSetOperatorRelBase extends PTransform<PCollectionList<Row>, PCo
@Override
public PCollection<Row> expand(PCollectionList<Row> inputs) {
- checkArgument(
- inputs.size() == 2,
- "Wrong number of arguments to %s: %s",
- beamRelNode.getClass().getSimpleName(),
- inputs);
- PCollection<Row> leftRows = inputs.get(0);
- PCollection<Row> rightRows = inputs.get(1);
+ Schema targetSchema = CalciteUtils.toSchema(beamRelNode.getRowType());
+
+ if (opType == OpType.UNION && all) {
+ PCollectionList<Row> coercedInputs = PCollectionList.empty(inputs.getPipeline());
+ for (PCollection<Row> 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<Row> 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<Row> expandPair(PCollection<Row> leftRows, PCollection<Row> rightRows) {
WindowFn leftWindow = leftRows.getWindowingStrategy().getWindowFn();
WindowFn rightWindow = rightRows.getWindowingStrategy().getWindowFn();
@@ -95,4 +121,54 @@ public class BeamSetOperatorRelBase extends PTransform<PCollectionList<Row>, PCo
lhsTag, rhsTag, opType, all)))
.setRowSchema(joined.getSchema().getField("key").getType().getRowSchema());
}
+
+ private static PCollection<Row> coerce(PCollection<Row> 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<Object> 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 c7a0b9136bb..a338d6b2e8b 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 class BeamSortRel extends Sort implements BeamRelNode {
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 class BeamSortRel extends Sort implements BeamRelNode {
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 e597b6c63ed..6fa4e383a2b 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.RexInputRef
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 class BeamWindowRel extends Window implements BeamRelNode {
List<Integer> 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 @@ public class BeamWindowRel extends Window implements BeamRelNode {
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<Integer> 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<Integer> partitionKeys;
@@ -208,8 +252,9 @@ public class BeamWindowRel extends Window implements BeamRelNode {
@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 class BeamWindowRel extends Window implements BeamRelNode {
*/
@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<PCollectionList<Row>, PCollection<Row>> {
@@ -309,16 +356,21 @@ public class BeamWindowRel extends Window implements BeamRelNode {
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 @@ public class BeamWindowRel extends Window implements BeamRelNode {
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.
+ *
+ * <p>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<Row> getRows(List<Row> input, int index) {
@@ -395,8 +472,9 @@ public class BeamWindowRel extends Window implements BeamRelNode {
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<Row> 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 e2df3c3a685..002477a717b 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 @@ package org.apache.beam.sdk.extensions.sql.impl.rule;
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 class BeamBasicAggregationRule extends RelOptRule {
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 f44839de062..156b3740e7b 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 class BeamCoGBKJoinRule extends RelOptRule {
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 class BeamCoGBKJoinRule extends RelOptRule {
public void onMatch(RelOptRuleCall call) {
Join join = (Join) call.rel(0);
+ RexNode condition = join.getCondition();
+ List<RexNode> conjuncts = new ArrayList<>();
+ RelOptUtil.decomposeConjunction(condition, conjuncts);
+
+ List<RexNode> equiConditions = new ArrayList<>();
+ List<RexNode> 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 class BeamCoGBKJoinRule extends RelOptRule {
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/BeamCoGBKJoinRule.java b/sdks/java/extensions/sql/src/main/java/org/apache/beam/sdk/extensions/sql/impl/rule/BeamCrossJoinRule.java
similarity index 67%
copy from sdks/java/extensions/sql/src/main/java/org/apache/beam/sdk/extensions/sql/impl/rule/BeamCoGBKJoinRule.java
copy to sdks/java/extensions/sql/src/main/java/org/apache/beam/sdk/extensions/sql/impl/rule/BeamCrossJoinRule.java
index f44839de062..18ceb9f0dd3 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/BeamCrossJoinRule.java
@@ -17,7 +17,7 @@
*/
package org.apache.beam.sdk.extensions.sql.impl.rule;
-import org.apache.beam.sdk.extensions.sql.impl.rel.BeamCoGBKJoinRel;
+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;
@@ -29,43 +29,40 @@ import org.apache.beam.vendor.calcite.v1_40_0.org.apache.calcite.rel.core.RelFac
import org.apache.beam.vendor.calcite.v1_40_0.org.apache.calcite.rel.logical.LogicalJoin;
/**
- * Rule to convert {@code LogicalJoin} node to {@code BeamCoGBKJoinRel} node.
- *
- * <p>This rule is matched when both the inputs to {@code LogicalJoin} node have the same
- * Boundedness i.e. either when both the inputs are {@code PCollection.IsBounded.BOUNDED} or when
- * both the inputs are {@code PCollection.IsBounded.UNBOUNDED}
- *
- * <p>As {@code BeamSideInputLookupJoinRel} also matches this condition when both the inputs are
- * {@code PCollection.IsBounded.BOUNDED}, to avoid conflicts, this rule is not matched when any of
- * the inputs to {@code LogicalJoin} node are Seekable.
+ * 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 BeamCoGBKJoinRule extends RelOptRule {
- public static final BeamCoGBKJoinRule INSTANCE = new BeamCoGBKJoinRule();
+public class BeamCrossJoinRule extends RelOptRule {
+ public static final BeamCrossJoinRule INSTANCE = new BeamCrossJoinRule();
- private BeamCoGBKJoinRule() {
+ private BeamCrossJoinRule() {
super(
operand(LogicalJoin.class, operand(RelNode.class, any()), operand(RelNode.class, any())),
RelFactories.LOGICAL_BUILDER,
- "BeamCoGBKJoinRule");
+ "BeamCrossJoinRule");
}
@Override
public boolean matches(RelOptRuleCall call) {
- // The Rule does not match when any of the inputs are Seekable
+ 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;
}
- PCollection.IsBounded boundednessOfLeftRel = BeamJoinRel.getBoundednessOfRelNode(call.rel(1));
- PCollection.IsBounded boundednessOfRightRel = BeamJoinRel.getBoundednessOfRelNode(call.rel(2));
- return (boundednessOfLeftRel == boundednessOfRightRel);
+ 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);
-
- BeamCoGBKJoinRel rel =
- new BeamCoGBKJoinRel(
+ BeamCrossJoinRel rel =
+ new BeamCrossJoinRel(
join.getCluster(),
join.getTraitSet().replace(BeamLogicalConvention.INSTANCE),
convert(
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 c442ae3d194..ad1d94bbc52 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 @@ public class BeamWindowRule extends ConverterRule {
@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 d9dbf881850..8eb23aef61f 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 final class BeamTableUtils {
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 fa2b6fe4adc..f79527e4da6 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 class BeamBuiltinAggregations {
static class FloatAvg extends Avg<Float> {
@Override
public @Nullable Float extractOutput(KV<Integer, BigDecimal> 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 class BeamBuiltinAggregations {
static class DoubleAvg extends Avg<Double> {
@Override
public @Nullable Double extractOutput(KV<Integer, BigDecimal> 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 413093d416d..b601c45f23e 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<T> extends Combine.CombineFn<T, Optional<T>, 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 <T> Combine.CombineFn<T, ?, T> 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 <T> Combine.CombineFn<T, ?, T> 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<T> {
+ private boolean seen;
+ private T value;
+
+ void set(T input) {
+ this.seen = true;
+ this.value = input;
+ }
+ }
+
+ private static class FirstValueCombineFn<T> extends Combine.CombineFn<T, ValueHolder<T>, T> {
private FirstValueCombineFn() {}
+ @Override
+ public ValueHolder<T> createAccumulator() {
+ return new ValueHolder<>();
+ }
+
+ @Override
+ public ValueHolder<T> addInput(ValueHolder<T> accumulator, T input) {
+ if (!accumulator.seen) {
+ accumulator.set(input);
+ }
+ return accumulator;
+ }
+
+ @Override
+ public ValueHolder<T> mergeAccumulators(Iterable<ValueHolder<T>> accumulators) {
+ throw new UnsupportedOperationException();
+ }
+
+ @Override
+ public T extractOutput(ValueHolder<T> accumulator) {
+ return accumulator.seen ? accumulator.value : null;
+ }
+ }
+
+ private static class LastValueCombineFn<T> extends Combine.CombineFn<T, ValueHolder<T>, T> {
+ private LastValueCombineFn() {}
+
+ @Override
+ public ValueHolder<T> createAccumulator() {
+ return new ValueHolder<>();
+ }
+
+ @Override
+ public ValueHolder<T> addInput(ValueHolder<T> accumulator, T input) {
+ accumulator.set(input);
+ return accumulator;
+ }
+
+ @Override
+ public ValueHolder<T> mergeAccumulators(Iterable<ValueHolder<T>> accumulators) {
+ throw new UnsupportedOperationException();
+ }
+
+ @Override
+ public T extractOutput(ValueHolder<T> accumulator) {
+ return accumulator.seen ? accumulator.value : null;
+ }
+ }
+
+ /** FIRST_VALUE(... IGNORE NULLS): the first non-null value in the frame. */
+ private static class FirstValueIgnoreNullsCombineFn<T>
+ extends Combine.CombineFn<T, Optional<T>, T> {
+ private FirstValueIgnoreNullsCombineFn() {}
+
@Override
public Optional<T> createAccumulator() {
return Optional.empty();
@@ -76,11 +174,10 @@ public class BeamBuiltinAnalyticFunctions {
@Override
public Optional<T> addInput(Optional<T> accumulator, T input) {
- Optional<T> 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 class BeamBuiltinAnalyticFunctions {
}
}
- private static class LastValueCombineFn<T> extends Combine.CombineFn<T, Optional<T>, T> {
- private LastValueCombineFn() {}
+ /** LAST_VALUE(... IGNORE NULLS): the last non-null value in the frame. */
+ private static class LastValueIgnoreNullsCombineFn<T>
+ extends Combine.CombineFn<T, Optional<T>, T> {
+ private LastValueIgnoreNullsCombineFn() {}
@Override
public Optional<T> createAccumulator() {
@@ -104,8 +203,10 @@ public class BeamBuiltinAnalyticFunctions {
@Override
public Optional<T> addInput(Optional<T> accumulator, T input) {
- Optional<T> r = Optional.of(input);
- return r;
+ if (input != null) {
+ return Optional.of(input);
+ }
+ return accumulator;
}
@Override
@@ -119,6 +220,74 @@ public class BeamBuiltinAnalyticFunctions {
}
}
+ /**
+ * 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<T>
+ extends PositionAwareCombineFn<BigDecimal, NthValueCombineFn.NthAccum<T>, T> {
+
+ /** Tracks how many eligible rows have been seen and the value captured at position n. */
+ private static class NthAccum<T> {
+ 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<T> createAccumulator() {
+ return new NthAccum<>();
+ }
+
+ @Override
+ public NthAccum<T> addInput(
+ NthAccum<T> accumulator,
+ BigDecimal input,
+ Long cursorWindow,
+ Long cursorPartition,
+ Long countPartition) {
+ return addInput(accumulator, input, input, cursorWindow, cursorPartition, countPartition);
+ }
+
+ @Override
+ @SuppressWarnings("unchecked")
+ public NthAccum<T> addInput(
+ NthAccum<T> 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<T> accumulator) {
+ return accumulator.captured ? accumulator.value : null;
+ }
+ }
+
// Numbering functions
public static <T> Combine.CombineFn<T, ?, T> numberingRowNumber() {
return new RowNumberCombineFn();
@@ -136,6 +305,10 @@ public class BeamBuiltinAnalyticFunctions {
return new PercentRankCombineFn();
}
+ public static <T> Combine.CombineFn<T, ?, T> numberingCumeDist() {
+ return new CumeDistCombineFn();
+ }
+
public abstract static class PositionAwareCombineFn<InputT, AccumT, OutputT>
extends Combine.CombineFn<InputT, AccumT, OutputT> {
public abstract AccumT addInput(
@@ -145,6 +318,23 @@ public class BeamBuiltinAnalyticFunctions {
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 class BeamBuiltinAnalyticFunctions {
}
}
+ /**
+ * 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)}.
+ *
+ * <p>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<T>
+ extends PositionAwareCombineFn<BigDecimal, KV<Long, Long>, Double> {
+
+ @Override
+ public KV<Long, Long> createAccumulator() {
+ return KV.of(0L, 0L);
+ }
+
+ @Override
+ public KV<Long, Long> addInput(
+ KV<Long, Long> accumulator,
+ BigDecimal input,
+ Long cursorPosition,
+ Long cursorPartition,
+ Long countPartition) {
+ return KV.of(accumulator.getKey() + 1L, countPartition);
+ }
+
+ @Override
+ public Double extractOutput(KV<Long, Long> accumulator) {
+ long total = accumulator.getValue();
+ if (total <= 0L) {
+ return 0.0;
+ }
+ return accumulator.getKey().doubleValue() / (double) total;
+ }
+ }
+
private static class PercentRankCombineFn<T>
extends PositionAwareCombineFn<BigDecimal, KV<Optional<Long>, KV<BigDecimal, Long>>, 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 906bac7add5..86f76a28eb7 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.schemas.Schema;
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 <em>Variance</em> on {@link Number} types.
@@ -68,7 +69,8 @@ import org.apache.beam.vendor.calcite.v1_40_0.org.apache.calcite.runtime.SqlFunc
@SuppressWarnings({
"rawtypes" // TODO(https://github.com/apache/beam/issues/20447)
})
-public class VarianceFn<T extends Number> extends Combine.CombineFn<T, VarianceAccumulator, T> {
+public class VarianceFn<T extends Number>
+ extends Combine.CombineFn<T, VarianceAccumulator, @Nullable T> {
static final MathContext MATH_CTX = new MathContext(10, RoundingMode.HALF_UP);
@@ -76,9 +78,6 @@ public class VarianceFn<T extends Number> extends Combine.CombineFn<T, VarianceA
private static final boolean POP = false;
private final boolean isSample; // flag to determine return value should be Variance Pop or Sample
- // When true, extractOutput returns the square root of the variance (i.e. standard deviation).
- // Beam's enumerable bridge cannot translate a SQRT call layered on top of a window VAR_SAMP, so
- // STDDEV_SAMP / STDDEV_POP are computed end-to-end inside this combiner instead.
private final boolean isStddev;
private final SerializableFunction<BigDecimal, T> decimalConverter;
@@ -146,9 +145,12 @@ public class VarianceFn<T extends Number> extends Combine.CombineFn<T, VarianceA
}
@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 class VarianceFn<T extends Number> extends Combine.CombineFn<T, VarianceA
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 d66fdb58d4a..a3719267030 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 class BeamAnalyticFunctionsTest extends BeamSqlDslBase {
pipeline.run();
}
+
+ @Test
+ public void testCumeDistFunction() throws Exception {
+ pipeline.enableAbandonedNodeEnforcement(false);
+ PCollection<Row> inputRows = inputData2();
+ String sql = "SELECT x, CUME_DIST() over (ORDER BY x ) as agg FROM PCOLLECTION";
+ PCollection<Row> result = inputRows.apply("sql", SqlTransform.query(sql));
+
+ Schema overResultSchema = Schema.builder().addInt32Field("x").addDoubleField("agg").build();
+
+ List<Row> 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<Row> inputRows = inputData2();
+ String sql = "SELECT x, COUNT(*) over () as agg FROM PCOLLECTION";
+ PCollection<Row> result = inputRows.apply("sql", SqlTransform.query(sql));
+
+ Schema overResultSchema = Schema.builder().addInt32Field("x").addInt64Field("agg").build();
+
+ List<Row> 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<Row> 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<Row> inputRows = inputDataTimestamp();
+ String sql =
+ "SELECT v, sum(v) over "
+ + "("
+ + "ORDER BY ts "
+ + "RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW"
+ + ")"
+ + " as total FROM PCOLLECTION";
+ PCollection<Row> result = inputRows.apply("sql", SqlTransform.query(sql));
+
+ Schema overResultSchema = Schema.builder().addInt32Field("v").addInt32Field("total").build();
+
+ List<Row> 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 b210bb0258f..8bae1107b61 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 class BeamSqlDslJoinTest {
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 9a0977e9018..27c9e14d874 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 class BeamCoGBKJoinRelBoundedVsBoundedTest extends BaseRelTest {
}
@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 class BeamCoGBKJoinRelBoundedVsBoundedTest extends BaseRelTest {
+ " 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<Row> 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 class BeamCoGBKJoinRelBoundedVsBoundedTest extends BaseRelTest {
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<Row> 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 f6bb0d50092..14a951d76cb 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 @@ package org.apache.beam.sdk.tpcds;
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 class QueryReader {
* @throws SqlParseException
*/
public static Set<String> 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 0d683c6b66d..cd5f8ad7988 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 class TableSchemaJSONLoader {
// 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 93cbc94f30a..8394acd70db 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 60072ea916c..a2a451b554e 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<String> 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<String> ALL_QUERY_NAMES = getAllQueryNames();
+ public static final List<String> 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 class TpcdsParametersReader {
}
/**
- * 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 class TpcdsParametersReader {
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<String> 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 763e8164dee..196e27e039c 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 class TpcdsSchemas {
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 class TpcdsSchemas {
.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 class TpcdsSchemas {
.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 class TpcdsSchemas {
.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 class TpcdsSchemas {
.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 class TpcdsSchemas {
.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 class TpcdsSchemas {
.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 class TpcdsSchemas {
.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 class TpcdsSchemas {
.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 class TpcdsSchemas {
.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 class TpcdsSchemas {
.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 class TpcdsSchemas {
.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 class TpcdsSchemas {
.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 class TpcdsSchemas {
.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 3eb84d32c57..1fd0422c814 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 class TableSchemaJSONLoaderTest {
+ "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 class TableSchemaJSONLoaderTest {
+ "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 class TableSchemaJSONLoaderTest {
+ "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 9b9c352541e..f49e4e1d1d9 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 class TpcdsParametersReaderTest {
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<String> 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 9f2397cf945..fdf9fea95aa 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 class TpcdsSchemasTest {
.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 class TpcdsSchemasTest {
.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 00000000000..e77f50c3b57
--- /dev/null
+++ b/sdks/python/apache_beam/testing/benchmarks/tpcds/README.md
@@ -0,0 +1,57 @@
+<!--
+ 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.
+-->
+
+# 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 00000000000..a0174df802d
--- /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 00000000000..8dac81f9449
--- /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<String>()
+ 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 00000000000..280226cd00e
--- /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 (<data_size>/<table>.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 00000000000..e7d45fc57d4
--- /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 00000000000..40831972a8e
--- /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 00000000000..b19da0f6268
--- /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 b8e7b855b99..b45c3f9ffa4 100644
--- a/settings.gradle.kts
+++ b/settings.gradle.kts
@@ -311,6 +311,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")