Remove redundancies from InputExtractor

Use Input class instead of TableEntry, and change behavior to simply return a List<Input> rather than a map.
prestodb · Aug 17, 2016 · b9abba0 · b9abba0
1 parent cb6b5f9
commit b9abba0
Show file tree

Hide file tree

Showing 2 changed files with 14 additions and 72 deletions.
diff --git a/presto-main/src/main/java/com/facebook/presto/execution/SqlQueryExecution.java b/presto-main/src/main/java/com/facebook/presto/execution/SqlQueryExecution.java
@@ -280,7 +280,7 @@ private PlanRoot doAnalyzeQuery()
         Plan plan = logicalPlanner.plan(analysis);
 
         // extract inputs
-        List<Input> inputs = new InputExtractor(metadata, stateMachine.getSession()).extract(plan.getRoot());
+        List<Input> inputs = new InputExtractor(metadata, stateMachine.getSession()).extractInputs(plan.getRoot());
         stateMachine.setInputs(inputs);
 
         // fragment the plan

diff --git a/presto-main/src/main/java/com/facebook/presto/sql/planner/InputExtractor.java b/presto-main/src/main/java/com/facebook/presto/sql/planner/InputExtractor.java
@@ -27,12 +27,10 @@
 import com.facebook.presto.sql.planner.plan.PlanVisitor;
 import com.facebook.presto.sql.planner.plan.TableScanNode;
 import com.google.common.collect.ImmutableList;
+import com.google.common.collect.ImmutableSet;
 
-import java.util.HashMap;
 import java.util.HashSet;
 import java.util.List;
-import java.util.Map;
-import java.util.Objects;
 import java.util.Optional;
 import java.util.Set;
 
@@ -47,39 +45,33 @@ public InputExtractor(Metadata metadata, Session session)
         this.session = session;
     }
 
-    public List<Input> extract(PlanNode root)
+    public List<Input> extractInputs(PlanNode root)
     {
         Visitor visitor = new Visitor();
         root.accept(visitor, null);
 
-        ImmutableList.Builder<Input> inputBuilder = ImmutableList.builder();
-        for (Map.Entry<TableEntry, Set<Column>> entry : visitor.getInputs().entrySet()) {
-            Input input = new Input(entry.getKey().getConnectorId(), entry.getKey().getSchema(), entry.getKey().getTable(), ImmutableList.copyOf(entry.getValue()));
-            inputBuilder.add(input);
-        }
-
-        return inputBuilder.build();
+        return ImmutableList.copyOf(visitor.getInputs());
     }
 
-    private static Column createColumnEntry(ColumnMetadata columnMetadata)
+    private static Column createColumn(ColumnMetadata columnMetadata)
     {
         return new Column(columnMetadata.getName(), columnMetadata.getType().toString());
     }
 
-    private static TableEntry createTableEntry(TableMetadata table)
+    private static Input createInput(TableMetadata table, Set<Column> columns)
     {
         SchemaTableName schemaTable = table.getTable();
-        return new TableEntry(table.getConnectorId(), schemaTable.getSchemaName(), schemaTable.getTableName());
+        return new Input(table.getConnectorId(), schemaTable.getSchemaName(), schemaTable.getTableName(), ImmutableList.copyOf(columns));
     }
 
     private class Visitor
             extends PlanVisitor<Void, Void>
     {
-        private final Map<TableEntry, Set<Column>> inputs = new HashMap<>();
+        private final ImmutableSet.Builder<Input> inputs = ImmutableSet.builder();
 
-        public Map<TableEntry, Set<Column>> getInputs()
+        public Set<Input> getInputs()
         {
-            return inputs;
+            return inputs.build();
         }
 
         @Override
@@ -91,11 +83,11 @@ public Void visitTableScan(TableScanNode node, Void context)
             Set<Column> columns = new HashSet<>();
             for (ColumnHandle columnHandle : node.getAssignments().values()) {
                 if (!columnHandle.equals(sampleWeightColumn.orElse(null))) {
-                    columns.add(createColumnEntry(metadata.getColumnMetadata(session, tableHandle, columnHandle)));
+                    columns.add(createColumn(metadata.getColumnMetadata(session, tableHandle, columnHandle)));
                 }
             }
 
-            inputs.put(createTableEntry(metadata.getTableMetadata(session, tableHandle)), columns);
+            inputs.add(createInput(metadata.getTableMetadata(session, tableHandle), columns));
 
             return null;
         }
@@ -109,11 +101,11 @@ public Void visitIndexSource(IndexSourceNode node, Void context)
             Set<Column> columns = new HashSet<>();
             for (ColumnHandle columnHandle : node.getAssignments().values()) {
                 if (!columnHandle.equals(sampleWeightColumn.orElse(null))) {
-                    columns.add(createColumnEntry(metadata.getColumnMetadata(session, tableHandle, columnHandle)));
+                    columns.add(createColumn(metadata.getColumnMetadata(session, tableHandle, columnHandle)));
                 }
             }
 
-            inputs.put(createTableEntry(metadata.getTableMetadata(session, tableHandle)), columns);
+            inputs.add(createInput(metadata.getTableMetadata(session, tableHandle), columns));
 
             return null;
         }
@@ -127,54 +119,4 @@ protected Void visitPlan(PlanNode node, Void context)
             return null;
         }
     }
-
-    private static final class TableEntry
-    {
-        private final String connectorId;
-        private final String schema;
-        private final String table;
-
-        private TableEntry(String connectorId, String schema, String table)
-        {
-            this.connectorId = connectorId;
-            this.schema = schema;
-            this.table = table;
-        }
-
-        public String getConnectorId()
-        {
-            return connectorId;
-        }
-
-        public String getSchema()
-        {
-            return schema;
-        }
-
-        public String getTable()
-        {
-            return table;
-        }
-
-        @Override
-        public int hashCode()
-        {
-            return Objects.hash(connectorId, schema, table);
-        }
-
-        @Override
-        public boolean equals(Object obj)
-        {
-            if (this == obj) {
-                return true;
-            }
-            if (obj == null || getClass() != obj.getClass()) {
-                return false;
-            }
-            final TableEntry other = (TableEntry) obj;
-            return Objects.equals(this.connectorId, other.connectorId) &&
-                    Objects.equals(this.schema, other.schema) &&
-                    Objects.equals(this.table, other.table);
-        }
-    }
 }