com.hazelcast.jet.sql.impl.opt.logical.WatermarkRules Maven / Gradle / Ivy

Go to download
Show more of this group Show more artifacts with this name
Show all versions of hazelcast-sql Show documentation
SQL Parser and Optimizer
There is a newer version: 5.5.0
/*
 * Copyright 2024 Hazelcast Inc.
 *
 * Licensed under the Hazelcast Community License (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 * http://hazelcast.com/hazelcast-community-license
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

package com.hazelcast.jet.sql.impl.opt.logical;

import com.hazelcast.shaded.com.google.common.collect.Iterables;
import com.hazelcast.jet.impl.util.Util;
import com.hazelcast.jet.sql.impl.aggregate.function.ImposeOrderFunction;
import com.hazelcast.jet.sql.impl.opt.OptUtils;
import com.hazelcast.jet.sql.impl.opt.metadata.WatermarkedFields;
import com.hazelcast.jet.sql.impl.opt.physical.visitor.RexToExpressionVisitor;
import com.hazelcast.sql.impl.QueryParameterMetadata;
import com.hazelcast.sql.impl.expression.Expression;
import com.hazelcast.shaded.org.apache.calcite.plan.Convention;
import com.hazelcast.shaded.org.apache.calcite.plan.HazelcastRelOptCluster;
import com.hazelcast.shaded.org.apache.calcite.plan.RelOptRule;
import com.hazelcast.shaded.org.apache.calcite.plan.RelOptRuleCall;
import com.hazelcast.shaded.org.apache.calcite.rel.logical.LogicalTableFunctionScan;
import com.hazelcast.shaded.org.apache.calcite.rex.RexCall;
import com.hazelcast.shaded.org.apache.calcite.rex.RexInputRef;
import com.hazelcast.shaded.org.apache.calcite.rex.RexNode;

import static com.hazelcast.jet.sql.impl.opt.metadata.HazelcastRelMdWatermarkedFields.watermarkedFieldByIndex;
import static com.hazelcast.sql.impl.plan.node.PlanNodeFieldTypeProvider.FAILING_FIELD_TYPE_PROVIDER;
import static com.hazelcast.shaded.org.apache.calcite.plan.RelOptRule.none;
import static com.hazelcast.shaded.org.apache.calcite.plan.RelOptRule.operand;
import static com.hazelcast.shaded.org.apache.calcite.plan.RelOptRule.operandJ;

final class WatermarkRules {

    /**
     * This rule converts IMPOSE_ORDER function call into
     * WatermarkLogicalRel. This rel is later pushed down into scan in
     * {@link #WATERMARK_INTO_SCAN_INSTANCE}.
     */
    @SuppressWarnings("AnonInnerLength")
    static final RelOptRule IMPOSE_ORDER_INSTANCE = new RelOptRule(
            operandJ(
                    LogicalTableFunctionScan.class,
                    Convention.NONE,
                    scan -> extractImposeOrderFunction(scan) != null,
                    none()
            ),
            WatermarkRules.class.getSimpleName() + "(Impose Order)"
    ) {
        @Override
        public void onMatch(RelOptRuleCall call) {
            LogicalTableFunctionScan scan = call.rel(0);

            int wmIndex = orderingColumnFieldIndex(scan);
            WatermarkLogicalRel wmRel = new WatermarkLogicalRel(
                    scan.getCluster(),
                    OptUtils.toLogicalConvention(scan.getTraitSet()),
                    Iterables.getOnlyElement(Util.toList(scan.getInputs(), OptUtils::toLogicalInput)),
                    lagExpression(scan),
                    wmIndex);

            if (wmIndex < 0) {
                call.transformTo(wmRel);
                return;
            }
            WatermarkedFields watermarkedFields = watermarkedFieldByIndex(wmIndex);
            if (watermarkedFields == null || watermarkedFields.isEmpty()) {
                call.transformTo(wmRel);
                return;
            }

            int watermarkedField = watermarkedFields.findFirst();

            DropLateItemsLogicalRel dropLateItemsRel = new DropLateItemsLogicalRel(
                    scan.getCluster(),
                    OptUtils.toLogicalConvention(scan.getTraitSet()),
                    wmRel,
                    watermarkedField
            );
            call.transformTo(dropLateItemsRel);
        }

        private int orderingColumnFieldIndex(LogicalTableFunctionScan function) {
            return ((RexInputRef) ((RexCall) extractOperand(function, 1)).getOperands().get(0)).getIndex();
        }

        private Expression lagExpression(LogicalTableFunctionScan function) {
            QueryParameterMetadata parameterMetadata = ((HazelcastRelOptCluster) function.getCluster()).getParameterMetadata();
            RexToExpressionVisitor visitor = new RexToExpressionVisitor(FAILING_FIELD_TYPE_PROVIDER, parameterMetadata);
            return extractOperand(function, 2).accept(visitor);
        }
    };

    /**
     * Push down {@link WatermarkLogicalRel} into {@link
     * FullScanLogicalRel}, if one is its input.
     */
    static final RelOptRule WATERMARK_INTO_SCAN_INSTANCE = new RelOptRule(
            operand(WatermarkLogicalRel.class, operand(FullScanLogicalRel.class, none())),
            WatermarkRules.class.getSimpleName() + "(Watermark Into Scan)"
    ) {
        @Override
        public void onMatch(RelOptRuleCall call) {
            WatermarkLogicalRel logicalWatermark = call.rel(0);
            FullScanLogicalRel logicalScan = call.rel(1);

            FullScanLogicalRel scan = new FullScanLogicalRel(
                    logicalWatermark.getCluster(),
                    logicalWatermark.getTraitSet(),
                    logicalScan.getTable(),
                    logicalWatermark.lagExpression(),
                    logicalWatermark.watermarkedColumnIndex()
            );
            call.transformTo(scan);
        }
    };

    private WatermarkRules() {
    }

    private static ImposeOrderFunction extractImposeOrderFunction(LogicalTableFunctionScan scan) {
        if (scan == null || !(scan.getCall() instanceof RexCall)) {
            return null;
        }
        RexCall call = (RexCall) scan.getCall();

        if (!(call.getOperator() instanceof ImposeOrderFunction)) {
            return null;
        }
        return (ImposeOrderFunction) call.getOperator();
    }

    private static RexNode extractOperand(LogicalTableFunctionScan function, int index) {
        return ((RexCall) function.getCall()).getOperands().get(index);
    }
}