Unflatten JDBC sources before planning (#93)

ryannedolan · web-flow · commit f2e4444621c2 · 2025-01-24T11:30:28.000-06:00
Integration tests not expected to pass.
diff --git a/hoptimator-k8s/src/main/java/com/linkedin/hoptimator/k8s/K8sDatabaseTable.java b/hoptimator-k8s/src/main/java/com/linkedin/hoptimator/k8s/K8sDatabaseTable.java
@@ -17,7 +17,7 @@
 import com.linkedin.hoptimator.k8s.models.V1alpha1Database;
 import com.linkedin.hoptimator.k8s.models.V1alpha1DatabaseList;
 import com.linkedin.hoptimator.k8s.models.V1alpha1DatabaseSpec;
-import com.linkedin.hoptimator.util.HoptimatorJdbcSchema;
+import com.linkedin.hoptimator.util.planner.HoptimatorJdbcSchema;
 
 
 public class K8sDatabaseTable extends K8sTable<V1alpha1Database, V1alpha1DatabaseList, K8sDatabaseTable.Row> {
diff --git a/hoptimator-k8s/src/main/java/com/linkedin/hoptimator/k8s/K8sEngineTable.java b/hoptimator-k8s/src/main/java/com/linkedin/hoptimator/k8s/K8sEngineTable.java
@@ -16,7 +16,7 @@
 import com.linkedin.hoptimator.k8s.models.V1alpha1Engine;
 import com.linkedin.hoptimator.k8s.models.V1alpha1EngineList;
 import com.linkedin.hoptimator.k8s.models.V1alpha1EngineSpec;
-import com.linkedin.hoptimator.util.HoptimatorJdbcSchema;
+import com.linkedin.hoptimator.util.planner.HoptimatorJdbcSchema;
 
 
 public class K8sEngineTable extends K8sTable<V1alpha1Engine, V1alpha1EngineList, K8sEngineTable.Row> {
diff --git a/hoptimator-util/src/main/java/com/linkedin/hoptimator/util/HoptimatorJdbcSchema.java b/hoptimator-util/src/main/java/com/linkedin/hoptimator/util/HoptimatorJdbcSchema.java
diff --git a/hoptimator-util/src/main/java/com/linkedin/hoptimator/util/planner/HoptimatorJdbcSchema.java b/hoptimator-util/src/main/java/com/linkedin/hoptimator/util/planner/HoptimatorJdbcSchema.java
@@ -0,0 +1,62 @@
+package com.linkedin.hoptimator.util.planner;
+
+import java.util.List;
+import javax.sql.DataSource;
+
+import org.apache.calcite.adapter.jdbc.JdbcSchema;
+import org.apache.calcite.adapter.jdbc.JdbcTable;
+import org.apache.calcite.linq4j.tree.Expression;
+import org.apache.calcite.schema.Schema;
+import org.apache.calcite.schema.Schemas;
+import org.apache.calcite.schema.SchemaPlus;
+import org.apache.calcite.schema.SchemaVersion;
+import org.apache.calcite.schema.Table;
+import org.apache.calcite.sql.SqlDialect;
+import org.apache.calcite.sql.SqlDialectFactory;
+import org.apache.calcite.sql.SqlDialectFactoryImpl;
+
+import com.linkedin.hoptimator.Database;
+import com.linkedin.hoptimator.Engine;
+import com.linkedin.hoptimator.util.planner.HoptimatorJdbcConvention;
+
+
+public class HoptimatorJdbcSchema extends JdbcSchema implements Database {
+
+  private final String database;
+  private final List<Engine> engines;
+  private final HoptimatorJdbcConvention convention;
+
+  public static HoptimatorJdbcSchema create(String database, String schema, DataSource dataSource,
+      SchemaPlus parentSchema, SqlDialect dialect, List<Engine> engines) {
+    Expression expression = Schemas.subSchemaExpression(parentSchema, schema, HoptimatorJdbcSchema.class);
+    HoptimatorJdbcConvention convention = new HoptimatorJdbcConvention(dialect, expression, database, engines);
+    return new HoptimatorJdbcSchema(database, schema, dataSource, dialect, convention, engines);
+  }
+
+  public HoptimatorJdbcSchema(String database, String schema, DataSource dataSource, SqlDialect dialect,
+      HoptimatorJdbcConvention convention, List<Engine> engines) {
+    super(dataSource, dialect, convention, null, schema);
+    this.database = database;
+    this.engines = engines;
+    this.convention = convention;
+  }
+
+  @Override
+  public String databaseName() {
+    return database;
+  }
+
+  public List<Engine> engines() {
+    return engines;
+  }
+
+  @Override
+  public Table getTable(String name) {
+    return new HoptimatorJdbcTable((JdbcTable) super.getTable(name), convention);
+  }
+
+  @Override
+  public Schema snapshot(SchemaVersion version) {
+    return this;
+  }
+}
diff --git a/hoptimator-util/src/main/java/com/linkedin/hoptimator/util/planner/HoptimatorJdbcTable.java b/hoptimator-util/src/main/java/com/linkedin/hoptimator/util/planner/HoptimatorJdbcTable.java
@@ -0,0 +1,76 @@
+package com.linkedin.hoptimator.util.planner;
+
+import java.util.Collection;
+import java.util.List;
+
+import com.linkedin.hoptimator.util.DataTypeUtils;
+
+import org.apache.calcite.adapter.java.AbstractQueryableTable;
+import org.apache.calcite.adapter.java.JavaTypeFactory;
+import org.apache.calcite.adapter.jdbc.JdbcTable;
+import org.apache.calcite.adapter.jdbc.JdbcTableScan;
+import org.apache.calcite.linq4j.QueryProvider;
+import org.apache.calcite.linq4j.Queryable;
+import org.apache.calcite.plan.RelOptCluster;
+import org.apache.calcite.plan.RelOptTable;
+import org.apache.calcite.prepare.Prepare.CatalogReader;
+import org.apache.calcite.rel.RelNode;
+import org.apache.calcite.rel.core.TableModify;
+import org.apache.calcite.rel.type.RelDataType;
+import org.apache.calcite.rel.type.RelDataTypeFactory;
+import org.apache.calcite.rex.RexNode;
+import org.apache.calcite.schema.ModifiableTable;
+import org.apache.calcite.schema.Table;
+import org.apache.calcite.schema.SchemaPlus;
+import org.apache.calcite.schema.TranslatableTable;
+import org.apache.calcite.schema.impl.AbstractTableQueryable;
+
+
+public class HoptimatorJdbcTable extends AbstractQueryableTable implements TranslatableTable,
+    ModifiableTable {
+
+  private final JdbcTable jdbcTable;
+  private final HoptimatorJdbcConvention convention;
+
+  public HoptimatorJdbcTable(JdbcTable jdbcTable, HoptimatorJdbcConvention convention) {
+    super(Object[].class);
+    this.jdbcTable = jdbcTable;
+    this.convention = convention;
+  }
+
+  public JdbcTable jdbcTable() {
+    return jdbcTable;
+  }
+
+  @Override
+  public RelDataType getRowType(RelDataTypeFactory factory) {
+    return DataTypeUtils.unflatten(jdbcTable.getRowType(factory), factory);
+  }
+
+  @Override
+  public RelNode toRel(RelOptTable.ToRelContext context,
+      RelOptTable relOptTable) {
+    return new HoptimatorJdbcTableScan(context.getCluster(), context.getTableHints(), relOptTable, this,
+        convention);
+  }
+
+  @Override
+  public TableModify toModificationRel(RelOptCluster cluster,
+      RelOptTable table, CatalogReader catalogReader, RelNode input,
+      TableModify.Operation operation, List<String> updateColumnList,
+      List<RexNode> sourceExpressionList, boolean flattened) {
+    return jdbcTable.toModificationRel(cluster, table, catalogReader, input, operation,
+        updateColumnList, sourceExpressionList, flattened);
+  }
+
+  @Override
+  public Collection getModifiableCollection() {
+    return null;
+  }
+
+  @Override
+  public <T> Queryable<T> asQueryable(QueryProvider queryProvider,
+      SchemaPlus schema, String tableName) {
+    return jdbcTable.asQueryable(queryProvider, schema, tableName);
+  }
+}
diff --git a/hoptimator-util/src/main/java/com/linkedin/hoptimator/util/planner/HoptimatorJdbcTableScan.java b/hoptimator-util/src/main/java/com/linkedin/hoptimator/util/planner/HoptimatorJdbcTableScan.java
@@ -0,0 +1,32 @@
+package com.linkedin.hoptimator.util.planner;
+
+import org.apache.calcite.adapter.jdbc.JdbcTable;
+import org.apache.calcite.adapter.jdbc.JdbcTableScan;
+import org.apache.calcite.plan.Convention;
+import org.apache.calcite.plan.RelOptCluster;
+import org.apache.calcite.plan.RelOptTable;
+import org.apache.calcite.plan.RelTraitSet;
+import org.apache.calcite.rel.RelNode;
+import org.apache.calcite.rel.core.TableScan;
+import org.apache.calcite.rel.hint.RelHint;
+
+import com.google.common.collect.ImmutableList;
+
+import java.util.List;
+
+import static org.apache.calcite.linq4j.Nullness.castNonNull;
+
+
+public class HoptimatorJdbcTableScan extends JdbcTableScan {
+  public final HoptimatorJdbcTable jdbcTable;
+
+  protected HoptimatorJdbcTableScan(
+      RelOptCluster cluster,
+      List<RelHint> hints,
+      RelOptTable table,
+      HoptimatorJdbcTable jdbcTable,
+      HoptimatorJdbcConvention jdbcConvention) {
+    super(cluster, hints, table, jdbcTable.jdbcTable(), jdbcConvention);
+    this.jdbcTable = jdbcTable;
+  }
+}
diff --git a/hoptimator-util/src/main/java/com/linkedin/hoptimator/util/planner/PipelineRules.java b/hoptimator-util/src/main/java/com/linkedin/hoptimator/util/planner/PipelineRules.java
@@ -89,8 +89,7 @@ static class PipelineTableScan extends TableScan implements PipelineRel {
     @Override
     public void implement(Implementor implementor) throws SQLException {
       RelDataTypeFactory typeFactory = new SqlTypeFactoryImpl(RelDataTypeSystem.DEFAULT);
-      implementor.addSource(database, table.getQualifiedName(),
-          DataTypeUtils.unflatten(table.getRowType(), typeFactory),
+      implementor.addSource(database, table.getQualifiedName(), table.getRowType(),
           Collections.emptyMap()); // TODO pass in table scan hints
     }
   }
diff --git a/hoptimator-util/src/main/java/com/linkedin/hoptimator/util/planner/ScriptImplementor.java b/hoptimator-util/src/main/java/com/linkedin/hoptimator/util/planner/ScriptImplementor.java
@@ -166,11 +166,12 @@ public QueryImplementor(RelNode relNode) {
     public void implement(SqlWriter w) {
       RelToSqlConverter converter = new RelToSqlConverter(w.getDialect());
       SqlImplementor.Result result = converter.visitRoot(relNode);
-      SqlSelect select = result.asSelect();
-      if (select.getSelectList() != null) {
+      SqlNode node = result.asStatement();
+      if (node instanceof SqlSelect && ((SqlSelect) node).getSelectList() != null) {
+        SqlSelect select = (SqlSelect) node;
         select.setSelectList((SqlNodeList) select.getSelectList().accept(REMOVE_ROW_CONSTRUCTOR));
       }
-      select.accept(new UnflattenMemberAccess(this)).unparse(w, 0, 0);
+      node.unparse(w, 0, 0);
     }
 
     // A `ROW(...)` operator which will unparse as just `(...)`.
@@ -189,32 +190,6 @@ public SqlNode visit(SqlCall call) {
         }
       }
     };
-
-    private static class UnflattenMemberAccess extends SqlShuttle {
-      private final Set<String> sinkFieldList;
-
-      UnflattenMemberAccess(QueryImplementor outer) {
-        this.sinkFieldList = outer.relNode.getRowType().getFieldList()
-            .stream()
-            .map(RelDataTypeField::getName)
-            .collect(Collectors.toSet());
-      }
-
-      // SqlShuttle gets called for every field in SELECT and every table name in FROM alike
-      // For fields in SELECT, we want to unflatten them as `FOO_BAR`, for tables `FOO.BAR`
-      @Override
-      public SqlNode visit(SqlIdentifier id) {
-        if (id.names.size() == 1 && sinkFieldList.contains(id.names.get(0))) {
-          id.assignNamesFrom(new SqlIdentifier(id.names.get(0).replaceAll("\\$", "_"), SqlParserPos.ZERO));
-        } else {
-          SqlIdentifier replacement = new SqlIdentifier(id.names.stream()
-              .flatMap(x -> Stream.of(x.split("\\$")))
-              .collect(Collectors.toList()), SqlParserPos.ZERO);
-          id.assignNamesFrom(replacement);
-        }
-        return id;
-      }
-    }
   }
 
   /**

Original file line number	Diff line number	Diff line change
`@@ -89,8 +89,7 @@ static class PipelineTableScan extends TableScan implements PipelineRel {`
`89`	`89`	`@Override`
`90`	`90`	`public void implement(Implementor implementor) throws SQLException {`
`91`	`91`	`RelDataTypeFactory typeFactory = new SqlTypeFactoryImpl(RelDataTypeSystem.DEFAULT);`
`92`		`- implementor.addSource(database, table.getQualifiedName(),`
`93`		`- DataTypeUtils.unflatten(table.getRowType(), typeFactory),`
	`92`	`+ implementor.addSource(database, table.getQualifiedName(), table.getRowType(),`
`94`	`93`	`Collections.emptyMap()); // TODO pass in table scan hints`
`95`	`94`	`}`
`96`	`95`	`}`