[ https://issues.apache.org/jira/browse/BEAM-8630?focusedWorklogId=361913&page=com.atlassian.jira.plugin.system.issuetabpanels:worklog-tabpanel#worklog-361913 ]
ASF GitHub Bot logged work on BEAM-8630: ---------------------------------------- Author: ASF GitHub Bot Created on: 20/Dec/19 22:20 Start Date: 20/Dec/19 22:20 Worklog Time Spent: 10m Work Description: apilloud commented on pull request #9913: [BEAM-8630] Prototype of BeamZetaSqlCalcRel URL: https://github.com/apache/beam/pull/9913#discussion_r360590462 ########## File path: sdks/java/extensions/sql/zetasql/src/main/java/org/apache/beam/sdk/extensions/sql/zetasql/BeamZetaSqlCalcRel.java ########## @@ -0,0 +1,259 @@ +/* + * 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.zetasql; + +import com.google.zetasql.AnalyzerOptions; +import com.google.zetasql.PreparedExpression; +import com.google.zetasql.Value; +import java.util.ArrayList; +import java.util.HashMap; +import java.util.List; +import java.util.Map; +import java.util.function.IntFunction; +import java.util.stream.Collectors; +import javax.annotation.Nullable; +import org.apache.beam.sdk.annotations.Internal; +import org.apache.beam.sdk.extensions.sql.impl.planner.BeamCostModel; +import org.apache.beam.sdk.extensions.sql.impl.planner.NodeStats; +import org.apache.beam.sdk.extensions.sql.impl.rel.BeamRelNode; +import org.apache.beam.sdk.extensions.sql.impl.rel.BeamSortRel; +import org.apache.beam.sdk.extensions.sql.impl.rel.BeamSqlRelUtils; +import org.apache.beam.sdk.extensions.sql.impl.rel.WithLimitableInput; +import org.apache.beam.sdk.extensions.sql.impl.utils.CalciteUtils; +import org.apache.beam.sdk.extensions.sql.meta.provider.bigquery.BeamBigQuerySqlDialect; +import org.apache.beam.sdk.schemas.Schema; +import org.apache.beam.sdk.schemas.Schema.Field; +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.values.PCollection; +import org.apache.beam.sdk.values.PCollectionList; +import org.apache.beam.sdk.values.Row; +import org.apache.beam.vendor.calcite.v1_20_0.org.apache.calcite.plan.RelOptCluster; +import org.apache.beam.vendor.calcite.v1_20_0.org.apache.calcite.plan.RelOptPlanner; +import org.apache.beam.vendor.calcite.v1_20_0.org.apache.calcite.plan.RelTraitSet; +import org.apache.beam.vendor.calcite.v1_20_0.org.apache.calcite.rel.RelNode; +import org.apache.beam.vendor.calcite.v1_20_0.org.apache.calcite.rel.core.Calc; +import org.apache.beam.vendor.calcite.v1_20_0.org.apache.calcite.rel.metadata.RelMetadataQuery; +import org.apache.beam.vendor.calcite.v1_20_0.org.apache.calcite.rel.rel2sql.SqlImplementor; +import org.apache.beam.vendor.calcite.v1_20_0.org.apache.calcite.rex.RexLocalRef; +import org.apache.beam.vendor.calcite.v1_20_0.org.apache.calcite.rex.RexNode; +import org.apache.beam.vendor.calcite.v1_20_0.org.apache.calcite.rex.RexProgram; +import org.apache.beam.vendor.calcite.v1_20_0.org.apache.calcite.sql.SqlDialect; +import org.apache.beam.vendor.calcite.v1_20_0.org.apache.calcite.sql.SqlIdentifier; +import org.apache.beam.vendor.calcite.v1_20_0.org.apache.calcite.sql.SqlNode; +import org.apache.beam.vendor.calcite.v1_20_0.org.apache.calcite.sql.parser.SqlParserPos; +import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.base.Preconditions; +import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableList; + +/** + * BeamRelNode to replace {@code Project} and {@code Filter} node based on the {@code ZetaSQL} + * expression evaluator. + */ +// TODO[BEAM-8630]: This class is currently a prototype and not used in runtime. +@Internal +public class BeamZetaSqlCalcRel extends Calc implements BeamRelNode, WithLimitableInput { + + private static final SqlDialect DIALECT = BeamBigQuerySqlDialect.DEFAULT; + private final SqlImplementor.Context context; + + public BeamZetaSqlCalcRel( + RelOptCluster cluster, RelTraitSet traits, RelNode input, RexProgram program) { + super(cluster, traits, input, program); + final IntFunction<SqlNode> fn = + i -> + new SqlIdentifier( + getProgram().getInputRowType().getFieldList().get(i).getName(), SqlParserPos.ZERO); + context = new SqlImplementor.SimpleContext(DIALECT, fn); + } + + @Override + public Calc copy(RelTraitSet traitSet, RelNode input, RexProgram program) { + return new BeamZetaSqlCalcRel(getCluster(), traitSet, input, program); + } + + @Override + public PTransform<PCollectionList<Row>, PCollection<Row>> buildPTransform() { + return new Transform(); + } + + private class Transform extends PTransform<PCollectionList<Row>, PCollection<Row>> { + @Override + public PCollection<Row> expand(PCollectionList<Row> pinput) { + Preconditions.checkArgument( + pinput.size() == 1, + "%s expected a single input PCollection, but received %d.", + BeamZetaSqlCalcRel.class.getSimpleName(), + pinput.size()); + PCollection<Row> upstream = pinput.get(0); + + final List<String> projects = + getProgram().getProjectList().stream() + .map(BeamZetaSqlCalcRel.this::unparseRexNode) + .collect(Collectors.toList()); + final RexNode condition = getProgram().getCondition(); + + // TODO[BEAM-8630]: validate sql expressions at pipeline construction time + Schema outputSchema = CalciteUtils.toSchema(getRowType()); + CalcFn calcFn = + new CalcFn( + projects, + condition == null ? null : unparseRexNode(condition), + upstream.getSchema(), + outputSchema); + return upstream.apply(ParDo.of(calcFn)).setRowSchema(outputSchema); + } + } + + private String unparseRexNode(RexNode rex) { + return context.toSql(getProgram(), rex).toSqlString(DIALECT).getSql(); + } + + @Override + public int getLimitCountOfSortRel() { + if (input instanceof BeamSortRel) { + return ((BeamSortRel) input).getCount(); + } + + throw new RuntimeException("Could not get the limit count from a non BeamSortRel input."); + } + + @Override + public NodeStats estimateNodeStats(RelMetadataQuery mq) { + NodeStats inputStat = BeamSqlRelUtils.getNodeStats(this.input, mq); + double selectivity = estimateFilterSelectivity(getInput(), program, mq); + + return inputStat.multiply(selectivity); + } + + private static double estimateFilterSelectivity( + RelNode child, RexProgram program, RelMetadataQuery mq) { + // Similar to calcite, if the calc node is representing filter operation we estimate the filter + // selectivity based on the number of equality conditions, number of inequality conditions, .... + RexLocalRef programCondition = program.getCondition(); + RexNode condition; + if (programCondition == null) { + condition = null; + } else { + condition = program.expandLocalRef(programCondition); + } + // Currently this gets the selectivity based on Calcite's Selectivity Handler (RelMdSelectivity) + return mq.getSelectivity(child, condition); + } + + @Override + public BeamCostModel beamComputeSelfCost(RelOptPlanner planner, RelMetadataQuery mq) { + NodeStats inputStat = BeamSqlRelUtils.getNodeStats(this.input, mq); + return BeamCostModel.FACTORY.makeCost(inputStat.getRowCount(), inputStat.getRate()); + } + + @Override + public boolean isInputSortRelAndLimitOnly() { + return (input instanceof BeamSortRel) && ((BeamSortRel) input).isLimitOnly(); + } + + /** + * {@code CalcFn} is the executor for a {@link BeamZetaSqlCalcRel} step. The implementation is + * based on the {@code ZetaSQL} expression evaluator. + */ + private static class CalcFn extends DoFn<Row, Row> { + private final List<String> projects; + @Nullable private final String condition; + private final Schema inputSchema; + private final Schema outputSchema; + private transient List<PreparedExpression> projectExps; + @Nullable private transient PreparedExpression conditionExp; + + CalcFn( + List<String> projects, + @Nullable String condition, + Schema inputSchema, + Schema outputSchema) { + Preconditions.checkArgument(projects.size() == outputSchema.getFieldCount()); + this.projects = ImmutableList.copyOf(projects); + this.condition = condition; + this.inputSchema = inputSchema; + this.outputSchema = outputSchema; + } + + @Setup + public void setup() { + AnalyzerOptions options = SqlAnalyzer.initAnalyzerOptions(); + for (Field field : inputSchema.getFields()) { + options.addExpressionColumn( + sanitize(field.getName()), ZetaSqlUtils.beamFieldTypeToZetaSqlType(field.getType())); + } + + projectExps = new ArrayList<>(); Review comment: `//TODO: Only use a single PreparedExpression for all condition and projects` ---------------------------------------------------------------- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. For queries about this service, please contact Infrastructure at: us...@infra.apache.org Issue Time Tracking ------------------- Worklog Id: (was: 361913) Time Spent: 1h 40m (was: 1.5h) > Prototype of BeamSQL Calc using ZetaSQL Expression Evaluator > ------------------------------------------------------------ > > Key: BEAM-8630 > URL: https://issues.apache.org/jira/browse/BEAM-8630 > Project: Beam > Issue Type: New Feature > Components: dsl-sql > Reporter: Yueyang Qiu > Assignee: Yueyang Qiu > Priority: Major > Time Spent: 1h 40m > Remaining Estimate: 0h > -- This message was sent by Atlassian Jira (v8.3.4#803005)