org.apache.asterix.optimizer.rules.IntroduceAutogenerateIDRule Maven / Gradle / Ivy

Go to download
/*
 * 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.asterix.optimizer.rules;

import java.util.ArrayList;
import java.util.List;

import org.apache.commons.lang3.mutable.Mutable;
import org.apache.commons.lang3.mutable.MutableObject;

import org.apache.asterix.aql.util.FunctionUtils;
import org.apache.asterix.metadata.declared.AqlDataSource;
import org.apache.asterix.metadata.declared.AqlDataSource.AqlDataSourceType;
import org.apache.asterix.metadata.declared.DatasetDataSource;
import org.apache.asterix.metadata.entities.InternalDatasetDetails;
import org.apache.asterix.om.base.AString;
import org.apache.asterix.om.constants.AsterixConstantValue;
import org.apache.asterix.om.functions.AsterixBuiltinFunctions;
import org.apache.hyracks.algebricks.common.exceptions.AlgebricksException;
import org.apache.hyracks.algebricks.core.algebra.base.ILogicalExpression;
import org.apache.hyracks.algebricks.core.algebra.base.ILogicalOperator;
import org.apache.hyracks.algebricks.core.algebra.base.IOptimizationContext;
import org.apache.hyracks.algebricks.core.algebra.base.LogicalOperatorTag;
import org.apache.hyracks.algebricks.core.algebra.base.LogicalVariable;
import org.apache.hyracks.algebricks.core.algebra.expressions.AbstractFunctionCallExpression;
import org.apache.hyracks.algebricks.core.algebra.expressions.ConstantExpression;
import org.apache.hyracks.algebricks.core.algebra.expressions.ScalarFunctionCallExpression;
import org.apache.hyracks.algebricks.core.algebra.expressions.VariableReferenceExpression;
import org.apache.hyracks.algebricks.core.algebra.operators.logical.AbstractLogicalOperator;
import org.apache.hyracks.algebricks.core.algebra.operators.logical.AssignOperator;
import org.apache.hyracks.algebricks.core.algebra.operators.logical.DataSourceScanOperator;
import org.apache.hyracks.algebricks.core.algebra.operators.logical.InsertDeleteOperator;
import org.apache.hyracks.algebricks.core.algebra.operators.logical.InsertDeleteOperator.Kind;
import org.apache.hyracks.algebricks.core.algebra.operators.logical.ProjectOperator;
import org.apache.hyracks.algebricks.core.algebra.operators.logical.visitors.VariableUtilities;
import org.apache.hyracks.algebricks.core.rewriter.base.IAlgebraicRewriteRule;

public class IntroduceAutogenerateIDRule implements IAlgebraicRewriteRule {

    @Override
    public boolean rewritePre(Mutable opRef, IOptimizationContext context) throws AlgebricksException {
        return false;
    }

    @Override
    public boolean rewritePost(Mutable opRef, IOptimizationContext context)
            throws AlgebricksException {

        // match: [insert to internal dataset with autogenerated id] - assign - project
        // produce: insert - assign - assign* - project
        // **
        // OR [insert to internal dataset with autogenerated id] - assign - [datasource scan]
        // produce insert - assign - assign* - datasource scan

        AbstractLogicalOperator currentOp = (AbstractLogicalOperator) opRef.getValue();
        if (currentOp.getOperatorTag() != LogicalOperatorTag.INSERT_DELETE) {
            return false;
        }

        InsertDeleteOperator insertOp = (InsertDeleteOperator) currentOp;
        if (insertOp.getOperation() != Kind.INSERT) {
            return false;
        }

        DatasetDataSource dds = (DatasetDataSource) insertOp.getDataSource();
        boolean autogenerated = ((InternalDatasetDetails) dds.getDataset().getDatasetDetails()).isAutogenerated();
        if (!autogenerated) {
            return false;
        }

        if (((AqlDataSource) insertOp.getDataSource()).getDatasourceType() != AqlDataSourceType.INTERNAL_DATASET) {
            return false;
        }

        AbstractLogicalOperator parentOp = (AbstractLogicalOperator) currentOp.getInputs().get(0).getValue();
        if (parentOp.getOperatorTag() != LogicalOperatorTag.ASSIGN) {
            return false;
        }
        AssignOperator assignOp = (AssignOperator) parentOp;
        LogicalVariable inputRecord;

        //bug here. will not work for internal datasets with filters since the pattern becomes [project-assign-assign-insert] <-this should be fixed->
        AbstractLogicalOperator grandparentOp = (AbstractLogicalOperator) parentOp.getInputs().get(0).getValue();
        if (grandparentOp.getOperatorTag() == LogicalOperatorTag.PROJECT) {
            ProjectOperator projectOp = (ProjectOperator) grandparentOp;
            inputRecord = projectOp.getVariables().get(0);
        } else if (grandparentOp.getOperatorTag() == LogicalOperatorTag.DATASOURCESCAN) {
            DataSourceScanOperator dssOp = (DataSourceScanOperator) grandparentOp;
            inputRecord = dssOp.getVariables().get(0);
        } else {
            return false;
        }

        List pkFieldName = ((InternalDatasetDetails) dds.getDataset().getDatasetDetails()).getPrimaryKey().get(
                0);
        ILogicalExpression rec0 = new VariableReferenceExpression(inputRecord);
        ILogicalExpression rec1 = createPrimaryKeyRecordExpression(pkFieldName);
        ILogicalExpression mergedRec = createRecordMergeFunction(rec0, rec1);
        ILogicalExpression nonNullMergedRec = createNotNullFunction(mergedRec);

        LogicalVariable v = context.newVar();
        AssignOperator newAssign = new AssignOperator(v, new MutableObject(nonNullMergedRec));
        newAssign.getInputs().add(new MutableObject(grandparentOp));
        assignOp.getInputs().set(0, new MutableObject(newAssign));
        VariableUtilities.substituteVariables(assignOp, inputRecord, v, context);
        VariableUtilities.substituteVariables(insertOp, inputRecord, v, context);
        context.computeAndSetTypeEnvironmentForOperator(newAssign);
        context.computeAndSetTypeEnvironmentForOperator(assignOp);
        context.computeAndSetTypeEnvironmentForOperator(insertOp);
        return true;
    }

    private ILogicalExpression createNotNullFunction(ILogicalExpression mergedRec) {
        List> args = new ArrayList<>();
        args.add(new MutableObject(mergedRec));
        AbstractFunctionCallExpression notNullFn = new ScalarFunctionCallExpression(
                FunctionUtils.getFunctionInfo(AsterixBuiltinFunctions.NOT_NULL), args);
        return notNullFn;
    }

    private AbstractFunctionCallExpression createPrimaryKeyRecordExpression(List pkFieldName) {
        //Create lowest level of nested uuid
        AbstractFunctionCallExpression uuidFn = new ScalarFunctionCallExpression(
                FunctionUtils.getFunctionInfo(AsterixBuiltinFunctions.CREATE_UUID));
        List> openRecordConsArgs = new ArrayList<>();
        Mutable pkFieldNameExpression = new MutableObject(
                new ConstantExpression(new AsterixConstantValue(new AString(pkFieldName.get(pkFieldName.size() - 1)))));
        openRecordConsArgs.add(pkFieldNameExpression);
        Mutable pkFieldValueExpression = new MutableObject(uuidFn);
        openRecordConsArgs.add(pkFieldValueExpression);
        AbstractFunctionCallExpression openRecFn = new ScalarFunctionCallExpression(
                FunctionUtils.getFunctionInfo(AsterixBuiltinFunctions.OPEN_RECORD_CONSTRUCTOR), openRecordConsArgs);

        //Create higher levels
        for (int i = pkFieldName.size() - 2; i > -1; i--) {
            AString fieldName = new AString(pkFieldName.get(i));
            openRecordConsArgs = new ArrayList<>();
            openRecordConsArgs.add(new MutableObject(new ConstantExpression(
                    new AsterixConstantValue(fieldName))));
            openRecordConsArgs.add(new MutableObject(openRecFn));
            openRecFn = new ScalarFunctionCallExpression(
                    FunctionUtils.getFunctionInfo(AsterixBuiltinFunctions.OPEN_RECORD_CONSTRUCTOR), openRecordConsArgs);
        }

        return openRecFn;
    }

    private AbstractFunctionCallExpression createRecordMergeFunction(ILogicalExpression rec0, ILogicalExpression rec1) {
        List> recordMergeFnArgs = new ArrayList<>();
        recordMergeFnArgs.add(new MutableObject<>(rec0));
        recordMergeFnArgs.add(new MutableObject<>(rec1));
        AbstractFunctionCallExpression recordMergeFn = new ScalarFunctionCallExpression(
                FunctionUtils.getFunctionInfo(AsterixBuiltinFunctions.RECORD_MERGE), recordMergeFnArgs);
        return recordMergeFn;
    }
}