Feat: TRAIN MODEL on join table

Author: taewhi

traindb-catalog/src/main/java/traindb/catalog/CatalogContext.java

@@ -142,6 +142,13 @@ void importSynopsis(String synopsisName, boolean isExternal, JSONObject exportMe
 
  void deleteTasks(Integer cnt) throws CatalogException;
 
+ /* Join */
+ public MTable createJoinTable(List<String> schemaNames, List<String> tableNames,
+ List<List<String>> columnNames, List<RelDataType> dataTypes,
+ String joinQuery) throws CatalogException;
+
+ public void dropJoinTable(String schemaName, String joinTableName) throws CatalogException;
+
  /* Common */
  void close();
 }

traindb-catalog/src/main/java/traindb/catalog/JDOCatalogContext.java

@@ -21,6 +21,7 @@
 import java.util.Comparator;
 import java.util.List;
 import java.util.Map;
+import java.util.UUID;
 import javax.jdo.PersistenceManager;
 import javax.jdo.Query;
 import javax.jdo.Transaction;
@@ -30,6 +31,7 @@
 import org.json.simple.JSONArray;
 import org.json.simple.JSONObject;
 import traindb.catalog.pm.MColumn;
+import traindb.catalog.pm.MJoin;
 import traindb.catalog.pm.MModel;
 import traindb.catalog.pm.MModeltype;
 import traindb.catalog.pm.MQueryLog;
@@ -149,12 +151,89 @@ private MTable addTable(String schemaName, String tableName, List<String> column
  return mTable;
  }
 
+ @Override
+ public MTable createJoinTable(List<String> schemaNames, List<String> tableNames,
+ List<List<String>> columnNames, List<RelDataType> dataTypes,
+ String joinQuery) throws CatalogException {
+ List<Long> srcTableIds = new ArrayList<>();
+ for (int i = 0; i < tableNames.size(); i++) {
+ MTable table = addTable(schemaNames.get(i), tableNames.get(i), columnNames.get(i),
+ dataTypes.get(i), "TABLE");
+ srcTableIds.add(table.getId());
+ }
+
+ UUID joinTableId = UUID.randomUUID();
+ String joinTableName = "__JOIN_" + joinTableId;
+ try {
+ // use first table's schema as join table's schema
+ MTable joinTable = new MTable(joinTableName, "JOIN", getSchema(schemaNames.get(0)));
+ pm.makePersistent(joinTable);
+ MTableExt joinTableExt = new MTableExt(joinTableName, "JOIN", joinQuery, joinTable);
+ pm.makePersistent(joinTableExt);
+
+ for (int i = 0; i < tableNames.size(); i++) {
+ List<String> colnames = columnNames.get(i);
+ RelDataType relDataType = dataTypes.get(i);
+ for (String colname : colnames) {
+ RelDataTypeField field = relDataType.getField(colname, true, false);
+ MColumn mColumn = new MColumn(field.getName(),
+ field.getType().getSqlTypeName().getJdbcOrdinal(),
+ field.getType().getPrecision(),
+ field.getType().getScale(),
+ field.getType().isNullable(),
+ joinTable);
+
+ pm.makePersistent(mColumn);
+ }
+ }
+
+ for (int i = 0; i < tableNames.size(); i++) {
+ MJoin join = new MJoin(joinTable.getId(), srcTableIds.get(i), columnNames.get(i));
+ pm.makePersistent(join);
+ }
+
+ return joinTable;
+ } catch (RuntimeException e) {
+ throw new CatalogException("failed to create join table", e);
+ }
+ }
+
+ @Override
+ public void dropJoinTable(String schemaName, String joinTableName) throws CatalogException {
+ Transaction tx = pm.currentTransaction();
+ try {
+ tx.begin();
+
+ MTable joinTable = getTable(schemaName, joinTableName);
+ if (joinTable == null) {
+ return;
+ }
+
+ Query query = pm.newQuery(MJoin.class);
+ setFilterPatterns(query, ImmutableMap.of("join_table_id", joinTable.getId()));
+ Collection<MJoin> mJoins = (List<MJoin>) query.execute();
+ pm.deletePersistentAll(mJoins);
+ pm.deletePersistent(joinTable);
+
+ tx.commit();
+ } catch (RuntimeException e) {
+ throw new CatalogException("failed to drop join table '" + joinTableName + "'", e);
+ } finally {
+ if (tx.isActive()) {
+ tx.rollback();
+ }
+ }
+ }
+
  @Override
  public MModel trainModel(
  String modeltypeName, String modelName, String schemaName, String tableName,
  List<String> columnNames, RelDataType dataType, @Nullable Long baseTableRows,
  @Nullable Long trainedRows, @Nullable String options) throws CatalogException {
- MTable mTable = addTable(schemaName, tableName, columnNames, dataType, "TABLE");
+ MTable mTable = getTable(schemaName, tableName);
+ if (mTable == null) {
+ mTable = addTable(schemaName, tableName, columnNames, dataType, "TABLE");
+ }
  try {
  MModeltype mModeltype = getModeltype(modeltypeName);
  MModel mModel = new MModel(
@@ -194,6 +273,7 @@ public void dropModel(String name) throws CatalogException {
  pm.deletePersistent(mTable);
  tx.commit();
  }
+ dropJoinTable(baseSchema, baseTable);
 
  Collection<MModel> baseSchemaModels = getModels(ImmutableMap.of("schema_name", baseSchema));
  if (baseSchemaModels == null || baseSchemaModels.size() == 0) {

traindb-catalog/src/main/java/traindb/catalog/pm/MJoin.java

@@ -0,0 +1,43 @@
+/*
+ * Licensed 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 traindb.catalog.pm;
+
+import java.util.List;
+import javax.jdo.annotations.PersistenceCapable;
+import javax.jdo.annotations.Persistent;
+
+@PersistenceCapable
+public final class MJoin {
+
+ @Persistent
+ private long join_table_id;
+
+ @Persistent
+ private long src_table_id;
+
+ @Persistent
+ private List<String> columns;
+
+ public MJoin(long joinTableId, long srcTableId, List<String> columns) {
+ this.join_table_id = joinTableId;
+ this.src_table_id = srcTableId;
+ this.columns = columns;
+ }
+
+ public List<String> getColumnNames() {
+ return columns;
+ }
+
+}

traindb-catalog/src/main/java/traindb/catalog/pm/MTable.java

@@ -56,6 +56,10 @@ public MTable(String name, String type, MSchema schema) {
  this.schema = schema;
  }
 
+ public long getId() {
+ return id;
+ }
+
  public String getTableName() {
  return table_name;
  }