Remove unnecessary property from PinotColumnMetadata

trinodb · Oct 22, 2022 · 3f5606b · 3f5606b
1 parent 4ded225
commit 3f5606b
Show file tree

Hide file tree

Showing 4 changed files with 23 additions and 32 deletions.
diff --git a/plugin/trino-pinot/src/main/java/io/trino/plugin/pinot/PinotColumnHandle.java b/plugin/trino-pinot/src/main/java/io/trino/plugin/pinot/PinotColumnHandle.java
@@ -15,7 +15,6 @@
 
 import com.fasterxml.jackson.annotation.JsonCreator;
 import com.fasterxml.jackson.annotation.JsonProperty;
-import com.google.common.collect.ImmutableMap;
 import io.trino.spi.connector.ColumnHandle;
 import io.trino.spi.connector.ColumnMetadata;
 import io.trino.spi.type.Type;
@@ -25,8 +24,6 @@
 
 import static com.google.common.base.MoreObjects.toStringHelper;
 import static com.google.common.base.Preconditions.checkState;
-import static io.trino.plugin.pinot.PinotMetadata.PINOT_COLUMN_NAME_PROPERTY;
-import static java.lang.String.format;
 import static java.util.Objects.requireNonNull;
 
 public class PinotColumnHandle
@@ -71,12 +68,6 @@ public PinotColumnHandle(
         this.pushedDownAggregateFunctionArgument = pushedDownAggregateFunctionArgument;
     }
 
-    public static PinotColumnHandle fromColumnMetadata(ColumnMetadata columnMetadata)
-    {
-        String columnName = (String) requireNonNull(columnMetadata.getProperties().get(PINOT_COLUMN_NAME_PROPERTY), format("Missing required column property '%s'", PINOT_COLUMN_NAME_PROPERTY));
-        return new PinotColumnHandle(columnName, columnMetadata.getType());
-    }
-
     @JsonProperty
     public String getColumnName()
     {
@@ -145,9 +136,6 @@ public ColumnMetadata getColumnMetadata()
         return ColumnMetadata.builder()
                 .setName(columnName)
                 .setType(dataType)
-                .setProperties(ImmutableMap.<String, Object>builder()
-                        .put(PINOT_COLUMN_NAME_PROPERTY, columnName)
-                        .buildOrThrow())
                 .build();
     }
 

diff --git a/plugin/trino-pinot/src/main/java/io/trino/plugin/pinot/PinotMetadata.java b/plugin/trino-pinot/src/main/java/io/trino/plugin/pinot/PinotMetadata.java
@@ -76,7 +76,6 @@
 import static com.google.common.collect.ImmutableMap.toImmutableMap;
 import static com.google.common.collect.Iterables.getOnlyElement;
 import static io.trino.collect.cache.SafeCaches.buildNonEvictableCache;
-import static io.trino.plugin.pinot.PinotColumnHandle.fromColumnMetadata;
 import static io.trino.plugin.pinot.PinotSessionProperties.isAggregationPushdownEnabled;
 import static io.trino.plugin.pinot.query.AggregateExpression.replaceIdentifier;
 import static io.trino.plugin.pinot.query.DynamicTablePqlExtractor.quoteIdentifier;
@@ -87,10 +86,9 @@
 public class PinotMetadata
         implements ConnectorMetadata
 {
-    public static final String PINOT_COLUMN_NAME_PROPERTY = "pinotColumnName";
     public static final String SCHEMA_NAME = "default";
 
-    private final NonEvictableLoadingCache<String, List<ColumnMetadata>> pinotTableColumnCache;
+    private final NonEvictableLoadingCache<String, List<PinotColumnHandle>> pinotTableColumnCache;
     private final int maxRowsPerBrokerQuery;
     private final AggregateFunctionRewriter<AggregateExpression, Void> aggregateFunctionRewriter;
     private final ImplementCountDistinct implementCountDistinct;
@@ -113,11 +111,11 @@ public PinotMetadata(
                 asyncReloading(new CacheLoader<>()
                 {
                     @Override
-                    public List<ColumnMetadata> load(String tableName)
+                    public List<PinotColumnHandle> load(String tableName)
                             throws Exception
                     {
                         Schema tablePinotSchema = pinotClient.getTableSchema(tableName);
-                        return getPinotColumnMetadataForPinotSchema(tablePinotSchema);
+                        return getPinotColumnHandlesForPinotSchema(tablePinotSchema);
                     }
                 }, executor));
 
@@ -199,9 +197,9 @@ public Map<String, ColumnHandle> getColumnHandles(ConnectorSession session, Conn
     public Map<String, ColumnHandle> getPinotColumnHandles(String tableName)
     {
         ImmutableMap.Builder<String, ColumnHandle> columnHandlesBuilder = ImmutableMap.builder();
-        for (ColumnMetadata columnMetadata : getColumnsMetadata(tableName)) {
-            columnHandlesBuilder.put(columnMetadata.getName(),
-                    fromColumnMetadata(columnMetadata));
+        String pinotTableName = pinotClient.getPinotTableNameFromTrinoTableName(tableName);
+        for (PinotColumnHandle columnHandle : getFromCache(pinotTableColumnCache, pinotTableName)) {
+            columnHandlesBuilder.put(columnHandle.getColumnName().toLowerCase(ENGLISH), columnHandle);
         }
         return columnHandlesBuilder.buildOrThrow();
     }
@@ -494,7 +492,9 @@ private static PinotColumnHandle resolveAggregateExpressionWithAlias(PinotColumn
     public List<ColumnMetadata> getColumnsMetadata(String tableName)
     {
         String pinotTableName = pinotClient.getPinotTableNameFromTrinoTableName(tableName);
-        return getFromCache(pinotTableColumnCache, pinotTableName);
+        return getFromCache(pinotTableColumnCache, pinotTableName).stream()
+                .map(PinotColumnHandle::getColumnMetadata)
+                .collect(toImmutableList());
     }
 
     private static <K, V> V getFromCache(LoadingCache<K, V> cache, K key)
@@ -526,17 +526,11 @@ private ConnectorTableMetadata getTableMetadata(SchemaTableName tableName)
         return new ConnectorTableMetadata(tableName, getColumnsMetadata(tableName.getTableName()));
     }
 
-    private List<ColumnMetadata> getPinotColumnMetadataForPinotSchema(Schema pinotTableSchema)
+    private List<PinotColumnHandle> getPinotColumnHandlesForPinotSchema(Schema pinotTableSchema)
     {
         return pinotTableSchema.getColumnNames().stream()
                 .filter(columnName -> !columnName.startsWith("$")) // Hidden columns starts with "$", ignore them as we can't use them in PQL
-                .map(columnName -> ColumnMetadata.builder()
-                        .setName(columnName)
-                        .setType(typeConverter.toTrinoType(pinotTableSchema.getFieldSpecFor(columnName)))
-                        .setProperties(ImmutableMap.<String, Object>builder()
-                                .put(PINOT_COLUMN_NAME_PROPERTY, columnName)
-                                .buildOrThrow())
-                        .build())
+                .map(columnName -> new PinotColumnHandle(columnName, typeConverter.toTrinoType(pinotTableSchema.getFieldSpecFor(columnName))))
                 .collect(toImmutableList());
     }
 

diff --git a/...ino-pinot/src/test/java/io/trino/plugin/pinot/BasePinotIntegrationConnectorSmokeTest.java b/...ino-pinot/src/test/java/io/trino/plugin/pinot/BasePinotIntegrationConnectorSmokeTest.java
@@ -1000,7 +1000,16 @@ public static Object of(
     @Override
     public void testShowCreateTable()
     {
-        assertQueryFails("SHOW CREATE TABLE region", "No PropertyMetadata for property: pinotColumnName");
+        assertThat((String) computeScalar("SHOW CREATE TABLE region"))
+                .isEqualTo(
+                        "CREATE TABLE %s.%s.region (\n" +
+                                "   regionkey bigint,\n" +
+                                "   updated_at_seconds bigint,\n" +
+                                "   name varchar,\n" +
+                                "   comment varchar\n" +
+                                ")",
+                        getSession().getCatalog().orElseThrow(),
+                        getSession().getSchema().orElseThrow());
     }
 
     @Override

diff --git a/plugin/trino-pinot/src/test/java/io/trino/plugin/pinot/TestPinotQueryBase.java b/plugin/trino-pinot/src/test/java/io/trino/plugin/pinot/TestPinotQueryBase.java
@@ -45,8 +45,8 @@ public class TestPinotQueryBase
 
     protected List<String> getColumnNames(String table)
     {
-        return pinotMetadata.getColumnsMetadata(table).stream()
-                .map(PinotColumnHandle::fromColumnMetadata)
+        return pinotMetadata.getPinotColumnHandles(table).values().stream()
+                .map(PinotColumnHandle.class::cast)
                 .map(PinotColumnHandle::getColumnName)
                 .collect(toImmutableList());
     }