[Coral-Incremental] Cost calculation for RelNode

coral-incremental/src/main/java/com/linkedin/coral/incremental/CostInfo.java

@@ -0,0 +1,17 @@
+/**
+ * Copyright 2024 LinkedIn Corporation. All rights reserved.
+ * Licensed under the BSD-2 Clause license.
+ * See LICENSE in the project root for license information.
+ */
+package com.linkedin.coral.incremental;
+
+public class CostInfo {
+  // TODO: we may also need to add TableName field.
+  Double cost;
+  Double row;
+
+  public CostInfo(Double cost, Double row) {
+    this.cost = cost;
+    this.row = row;
+  }
+}

coral-incremental/src/main/java/com/linkedin/coral/incremental/RelNodeCostEstimator.java

@@ -0,0 +1,168 @@
+/**
+ * Copyright 2024 LinkedIn Corporation. All rights reserved.
+ * Licensed under the BSD-2 Clause license.
+ * See LICENSE in the project root for license information.
+ */
+package com.linkedin.coral.incremental;
+
+import java.util.ArrayList;
+import java.util.HashMap;
+import java.util.Iterator;
+import java.util.List;
+import java.util.Map;
+
+import org.apache.calcite.plan.RelOptTable;
+import org.apache.calcite.rel.RelNode;
+import org.apache.calcite.rel.core.TableScan;
+import org.apache.calcite.rel.logical.LogicalJoin;
+import org.apache.calcite.rel.logical.LogicalProject;
+import org.apache.calcite.rel.logical.LogicalUnion;
+import org.apache.calcite.rel.type.RelDataType;
+import org.apache.calcite.rel.type.RelDataTypeField;
+import org.apache.calcite.rex.RexCall;
+import org.apache.calcite.rex.RexInputRef;
+import org.apache.calcite.rex.RexNode;
+
+import static java.lang.Math.*;
+
+
+public class RelNodeCostEstimator {
+
+  class JoinKey {
+    String leftTableName;
+    String rightTableName;
+    String leftFieldName;
+    String rightFieldName;
+
+    public JoinKey(String leftTableName, String rightTableName, String leftFieldName, String rightFieldName) {
+      this.leftTableName = leftTableName;
+      this.rightTableName = rightTableName;
+      this.leftFieldName = leftFieldName;
+      this.rightFieldName = rightFieldName;
+    }
+  }
+
+  private Map<String, Double> stat = new HashMap<>();
+
+  public void setStat(Map<String, Double> stat) {
+    this.stat = stat;
+  }
+
+  private Double IOCostParam = 1.0;
+
+  private Double shuffleCostParam = 1.0;
+
+  public Double getCost(RelNode rel) {
+    CostInfo executionCostInfo = getExecutionCost(rel);
+    Double IOCost = executionCostInfo.row * IOCostParam;
+    return executionCostInfo.cost * shuffleCostParam + IOCost;
+  }
+
+  public CostInfo getExecutionCost(RelNode rel) {
+    if (rel instanceof TableScan) {
+      return getExecutionCostTableScan((TableScan) rel);
+    } else if (rel instanceof LogicalJoin) {
+      return getExecutionCostJoin((LogicalJoin) rel);
+    } else if (rel instanceof LogicalUnion) {
+      return getExecutionCostUnion((LogicalUnion) rel);
+    } else if (rel instanceof LogicalProject) {
+      return getExecutionCostProject((LogicalProject) rel);
+    }
+    return new CostInfo(0.0, 0.0);
+  }
+
+  private CostInfo getExecutionCostTableScan(TableScan scan) {
+    RelOptTable originalTable = scan.getTable();
+    String tableName = getTableName(originalTable);
+    Double row = stat.getOrDefault(tableName, 5.0);
+    return new CostInfo(row, row);
+  }
+
+  private String getTableName(RelOptTable table) {
+    return String.join(".", table.getQualifiedName());
+  }
+
+  private CostInfo getExecutionCostJoin(LogicalJoin join) {
+    RelNode left = join.getLeft();
+    RelNode right = join.getRight();
+    //    if (!(left instanceof TableScan) || !(right instanceof TableScan))
+    //    {
+    //      return new CostInfo(0.0, 0.0);
+    //    }
+    CostInfo leftCost = getExecutionCost(left);
+    CostInfo rightCost = getExecutionCost(right);
+    Double joinSize = estimateJoinSize(join, leftCost.row, rightCost.row);
+    return new CostInfo(max(leftCost.cost, rightCost.cost), joinSize);
+  }
+
+  private List<JoinKey> findJoinKeys(LogicalJoin join) {
+    List<JoinKey> joinKeys = new ArrayList<>();
+    RexNode condition = join.getCondition();
+    if (condition instanceof RexCall) {
+      processRexCall((RexCall) condition, join, joinKeys);
+    }
+    return joinKeys;
+  }
+
+  private void processRexCall(RexCall call, LogicalJoin join, List<JoinKey> joinKeys) {
+    if (call.getOperator().getName().equalsIgnoreCase("AND")) {
+      // Process each operand of the AND separately
+      for (RexNode operand : call.getOperands()) {
+        if (operand instanceof RexCall) {
+          processRexCall((RexCall) operand, join, joinKeys);
+        }
+      }
+    } else {
+      // Process the join condition (e.g., EQUALS)
+      List<RexNode> operands = call.getOperands();
+      if (operands.size() == 2 && operands.get(0) instanceof RexInputRef && operands.get(1) instanceof RexInputRef) {
+        RexInputRef leftRef = (RexInputRef) operands.get(0);
+        RexInputRef rightRef = (RexInputRef) operands.get(1);
+        RelDataType leftType = join.getLeft().getRowType();
+        RelDataType rightType = join.getRight().getRowType();
+
+        int leftIndex = leftRef.getIndex();
+        int rightIndex = rightRef.getIndex() - leftType.getFieldCount();
+
+        RelDataTypeField leftField = leftType.getFieldList().get(leftIndex);
+        String leftTableName = getTableName(join.getLeft().getTable());
+        String leftFieldName = leftField.getName();
+        RelDataTypeField rightField = rightType.getFieldList().get(rightIndex);
+        String rightTableName = getTableName(join.getRight().getTable());
+        String rightFieldName = rightField.getName();
+
+        joinKeys.add(new JoinKey(leftTableName, rightTableName, leftFieldName, rightFieldName));
+      }
+    }
+  }
+
+  private Double estimateJoinSelectivity(List<JoinKey> joinKeys) {
+    if (joinKeys.size() == 1 && joinKeys.get(0).leftFieldName == "x") {
+      return 0.1;
+    }
+    return 1.0;
+  }
+
+  private Double estimateJoinSize(LogicalJoin join, Double leftSize, Double rightSize) {
+    Double selectivity = estimateJoinSelectivity(findJoinKeys(join));
+    return leftSize * rightSize * selectivity;
+  }
+
+  private CostInfo getExecutionCostUnion(LogicalUnion union) {
+    Double unionCost = 0.0;
+    Double unionSize = 0.0;
+    RelNode input;
+    for (Iterator var4 = union.getInputs().iterator(); var4.hasNext();) {
+      input = (RelNode) var4.next();
+      CostInfo inputCost = getExecutionCost(input);
+      unionSize += inputCost.row;
+      unionCost = max(inputCost.cost, unionCost);
+    }
+    unionCost *= 1.5;
+    return new CostInfo(unionCost, unionSize);
+  }
+
+  private CostInfo getExecutionCostProject(LogicalProject project) {
+    return getExecutionCost(project.getInput());
+  }
+}

coral-incremental/src/main/java/com/linkedin/coral/incremental/RelNodeGenerationTransformer.java

@@ -0,0 +1,167 @@
+/**
+ * Copyright 2024 LinkedIn Corporation. All rights reserved.
+ * Licensed under the BSD-2 Clause license.
+ * See LICENSE in the project root for license information.
+ */
+package com.linkedin.coral.incremental;
+
+import java.util.ArrayList;
+import java.util.Arrays;
+import java.util.List;
+import java.util.stream.Collectors;
+import java.util.stream.IntStream;
+
+import org.apache.calcite.plan.RelOptTable;
+import org.apache.calcite.prepare.RelOptTableImpl;
+import org.apache.calcite.rel.RelNode;
+import org.apache.calcite.rel.RelShuttle;
+import org.apache.calcite.rel.RelShuttleImpl;
+import org.apache.calcite.rel.core.TableScan;
+import org.apache.calcite.rel.logical.LogicalAggregate;
+import org.apache.calcite.rel.logical.LogicalFilter;
+import org.apache.calcite.rel.logical.LogicalJoin;
+import org.apache.calcite.rel.logical.LogicalProject;
+import org.apache.calcite.rel.logical.LogicalTableScan;
+import org.apache.calcite.rel.logical.LogicalUnion;
+import org.apache.calcite.rex.RexBuilder;
+import org.apache.calcite.rex.RexNode;
+
+
+public class RelNodeGenerationTransformer {
+
+  static RelNode convertRelPrev(RelNode originalNode) {
+    RelShuttle converter = new RelShuttleImpl() {
+      @Override
+      public RelNode visit(TableScan scan) {
+        RelOptTable originalTable = scan.getTable();
+        List<String> incrementalNames = new ArrayList<>(originalTable.getQualifiedName());
+        String deltaTableName = incrementalNames.remove(incrementalNames.size() - 1) + "_prev";
+        incrementalNames.add(deltaTableName);
+        RelOptTable incrementalTable =
+            RelOptTableImpl.create(originalTable.getRelOptSchema(), originalTable.getRowType(), incrementalNames, null);
+        return LogicalTableScan.create(scan.getCluster(), incrementalTable);
+      }
+
+      @Override
+      public RelNode visit(LogicalJoin join) {
+        RelNode left = join.getLeft();
+        RelNode right = join.getRight();
+        RelNode prevLeft = convertRelPrev(left);
+        RelNode prevRight = convertRelPrev(right);
+        RexBuilder rexBuilder = join.getCluster().getRexBuilder();
+
+        LogicalProject p3 = createProjectOverJoin(join, prevLeft, prevRight, rexBuilder);
+
+        return p3;
+      }
+
+      @Override
+      public RelNode visit(LogicalFilter filter) {
+        RelNode transformedChild = convertRelPrev(filter.getInput());
+
+        return LogicalFilter.create(transformedChild, filter.getCondition());
+      }
+
+      @Override
+      public RelNode visit(LogicalProject project) {
+        RelNode transformedChild = convertRelPrev(project.getInput());
+        return LogicalProject.create(transformedChild, project.getProjects(), project.getRowType());
+      }
+
+      @Override
+      public RelNode visit(LogicalUnion union) {
+        List<RelNode> children = union.getInputs();
+        List<RelNode> transformedChildren =
+            children.stream().map(child -> convertRelPrev(child)).collect(Collectors.toList());
+        return LogicalUnion.create(transformedChildren, union.all);
+      }
+
+      @Override
+      public RelNode visit(LogicalAggregate aggregate) {
+        RelNode transformedChild = convertRelPrev(aggregate.getInput());
+        return LogicalAggregate.create(transformedChild, aggregate.getGroupSet(), aggregate.getGroupSets(),
+            aggregate.getAggCallList());
+      }
+    };
+    return originalNode.accept(converter);
+  }
+
+  private RelNodeGenerationTransformer() {
+  }
+
+  public static RelNode convertRelIncremental(RelNode originalNode) {
+    RelShuttle converter = new RelShuttleImpl() {
+      @Override
+      public RelNode visit(TableScan scan) {
+        RelOptTable originalTable = scan.getTable();
+        List<String> incrementalNames = new ArrayList<>(originalTable.getQualifiedName());
+        String deltaTableName = incrementalNames.remove(incrementalNames.size() - 1) + "_delta";
+        incrementalNames.add(deltaTableName);
+        RelOptTable incrementalTable =
+            RelOptTableImpl.create(originalTable.getRelOptSchema(), originalTable.getRowType(), incrementalNames, null);
+        return LogicalTableScan.create(scan.getCluster(), incrementalTable);
+      }
+
+      @Override
+      public RelNode visit(LogicalJoin join) {
+        RelNode left = join.getLeft();
+        RelNode right = join.getRight();
+        RelNode prevLeft = convertRelPrev(left);
+        RelNode prevRight = convertRelPrev(right);
+        RelNode incrementalLeft = convertRelIncremental(left);
+        RelNode incrementalRight = convertRelIncremental(right);
+
+        RexBuilder rexBuilder = join.getCluster().getRexBuilder();
+
+        LogicalProject p1 = createProjectOverJoin(join, prevLeft, incrementalRight, rexBuilder);
+        LogicalProject p2 = createProjectOverJoin(join, incrementalLeft, prevRight, rexBuilder);
+        LogicalProject p3 = createProjectOverJoin(join, incrementalLeft, incrementalRight, rexBuilder);
+
+        LogicalUnion unionAllJoins =
+            LogicalUnion.create(Arrays.asList(LogicalUnion.create(Arrays.asList(p1, p2), true), p3), true);
+        return unionAllJoins;
+      }
+
+      @Override
+      public RelNode visit(LogicalFilter filter) {
+        RelNode transformedChild = convertRelIncremental(filter.getInput());
+        return LogicalFilter.create(transformedChild, filter.getCondition());
+      }
+
+      @Override
+      public RelNode visit(LogicalProject project) {
+        RelNode transformedChild = convertRelIncremental(project.getInput());
+        return LogicalProject.create(transformedChild, project.getProjects(), project.getRowType());
+      }
+
+      @Override
+      public RelNode visit(LogicalUnion union) {
+        List<RelNode> children = union.getInputs();
+        List<RelNode> transformedChildren =
+            children.stream().map(child -> convertRelIncremental(child)).collect(Collectors.toList());
+        return LogicalUnion.create(transformedChildren, union.all);
+      }
+
+      @Override
+      public RelNode visit(LogicalAggregate aggregate) {
+        RelNode transformedChild = convertRelIncremental(aggregate.getInput());
+        return LogicalAggregate.create(transformedChild, aggregate.getGroupSet(), aggregate.getGroupSets(),
+            aggregate.getAggCallList());
+      }
+    };
+    return originalNode.accept(converter);
+  }
+
+  private static LogicalProject createProjectOverJoin(LogicalJoin join, RelNode left, RelNode right,
+      RexBuilder rexBuilder) {
+    LogicalJoin incrementalJoin =
+        LogicalJoin.create(left, right, join.getCondition(), join.getVariablesSet(), join.getJoinType());
+    ArrayList<RexNode> projects = new ArrayList<>();
+    ArrayList<String> names = new ArrayList<>();
+    IntStream.range(0, incrementalJoin.getRowType().getFieldList().size()).forEach(i -> {
+      projects.add(rexBuilder.makeInputRef(incrementalJoin, i));
+      names.add(incrementalJoin.getRowType().getFieldNames().get(i));
+    });
+    return LogicalProject.create(incrementalJoin, projects, names);
+  }
+}

coral-incremental/src/test/java/com/linkedin/coral/incremental/RelNodeCostEstimatorTest.java

@@ -0,0 +1,10 @@
+/**
+ * Copyright 2024 LinkedIn Corporation. All rights reserved.
+ * Licensed under the BSD-2 Clause license.
+ * See LICENSE in the project root for license information.
+ */
+package com.linkedin.coral.incremental;
+
+public class RelNodeCostEstimatorTest {
+
+}