Fix type of DeltaLakeTableHandle.projectedColumns

It's a collection of Delta column handles.
trinodb · May 10, 2023 · e7f11d3 · e7f11d3
1 parent 88a16c4
commit e7f11d3
Show file tree

Hide file tree

Showing 4 changed files with 15 additions and 17 deletions.
diff --git a/plugin/trino-delta-lake/src/main/java/io/trino/plugin/deltalake/DeltaLakeMetadata.java b/plugin/trino-delta-lake/src/main/java/io/trino/plugin/deltalake/DeltaLakeMetadata.java
@@ -567,9 +567,7 @@ public Map<String, ColumnHandle> getColumnHandles(ConnectorSession session, Conn
     {
         DeltaLakeTableHandle table = checkValidTableHandle(tableHandle);
         return table.getProjectedColumns()
-                .map(projectedColumns -> (List<DeltaLakeColumnHandle>) projectedColumns.stream()
-                        .map(DeltaLakeColumnHandle.class::cast) // TODO DeltaLakeTableHandle.projectedColumns should be a collection of DeltaLakeColumnHandle
-                        .collect(toImmutableList()))
+                .map(projectColumns -> (Collection<DeltaLakeColumnHandle>) projectColumns)
                 .orElseGet(() -> getColumns(table.getMetadataEntry())).stream()
                 // This method does not calculate column name for the projected columns
                 .peek(handle -> checkArgument(handle.isBaseColumn(), "Unsupported projected column: %s", handle))
@@ -2283,7 +2281,9 @@ public Optional<ProjectionApplicationResult<ConnectorTableHandle>> applyProjecti
         // all references are simple variables
         if (!isProjectionPushdownEnabled(session)
                 || columnProjections.values().stream().allMatch(ProjectedColumnRepresentation::isVariable)) {
-            Set<ColumnHandle> projectedColumns = ImmutableSet.copyOf(assignments.values());
+            Set<DeltaLakeColumnHandle> projectedColumns = assignments.values().stream()
+                    .map(DeltaLakeColumnHandle.class::cast)
+                    .collect(toImmutableSet());
             // Check if column was projected already in previous call
             if (deltaLakeTableHandle.getProjectedColumns().isPresent()
                     && deltaLakeTableHandle.getProjectedColumns().get().equals(projectedColumns)) {
@@ -2306,7 +2306,7 @@ public Optional<ProjectionApplicationResult<ConnectorTableHandle>> applyProjecti
 
         Map<String, Assignment> newAssignments = new HashMap<>();
         ImmutableMap.Builder<ConnectorExpression, Variable> newVariablesBuilder = ImmutableMap.builder();
-        ImmutableSet.Builder<ColumnHandle> projectedColumnsBuilder = ImmutableSet.builder();
+        ImmutableSet.Builder<DeltaLakeColumnHandle> projectedColumnsBuilder = ImmutableSet.builder();
 
         for (Map.Entry<ConnectorExpression, ProjectedColumnRepresentation> entry : columnProjections.entrySet()) {
             ConnectorExpression expression = entry.getKey();

diff --git a/plugin/trino-delta-lake/src/main/java/io/trino/plugin/deltalake/DeltaLakeSplitManager.java b/plugin/trino-delta-lake/src/main/java/io/trino/plugin/deltalake/DeltaLakeSplitManager.java
@@ -226,7 +226,7 @@ private static boolean mayAnyDataColumnProjected(DeltaLakeTableHandle tableHandl
             return true;
         }
         return tableHandle.getProjectedColumns().get().stream()
-                .map(columnHandle -> ((DeltaLakeColumnHandle) columnHandle).getColumnType())
+                .map(DeltaLakeColumnHandle::getColumnType)
                 .anyMatch(DeltaLakeColumnType.REGULAR::equals);
     }
 

diff --git a/plugin/trino-delta-lake/src/main/java/io/trino/plugin/deltalake/DeltaLakeTableHandle.java b/plugin/trino-delta-lake/src/main/java/io/trino/plugin/deltalake/DeltaLakeTableHandle.java
@@ -18,7 +18,6 @@
 import com.fasterxml.jackson.annotation.JsonProperty;
 import io.airlift.units.DataSize;
 import io.trino.plugin.deltalake.transactionlog.MetadataEntry;
-import io.trino.spi.connector.ColumnHandle;
 import io.trino.spi.connector.SchemaTableName;
 import io.trino.spi.predicate.TupleDomain;
 
@@ -51,7 +50,7 @@ public enum WriteType
     private final Optional<WriteType> writeType;
     private final long readVersion;
 
-    private final Optional<Set<ColumnHandle>> projectedColumns;
+    private final Optional<Set<DeltaLakeColumnHandle>> projectedColumns;
     // UPDATE only: The list of columns being updated
     private final Optional<List<DeltaLakeColumnHandle>> updatedColumns;
     // UPDATE only: The list of columns which need to be copied when applying updates to the new Parquet file
@@ -74,7 +73,7 @@ public DeltaLakeTableHandle(
             @JsonProperty("enforcedPartitionConstraint") TupleDomain<DeltaLakeColumnHandle> enforcedPartitionConstraint,
             @JsonProperty("nonPartitionConstraint") TupleDomain<DeltaLakeColumnHandle> nonPartitionConstraint,
             @JsonProperty("writeType") Optional<WriteType> writeType,
-            @JsonProperty("projectedColumns") Optional<Set<ColumnHandle>> projectedColumns,
+            @JsonProperty("projectedColumns") Optional<Set<DeltaLakeColumnHandle>> projectedColumns,
             @JsonProperty("updatedColumns") Optional<List<DeltaLakeColumnHandle>> updatedColumns,
             @JsonProperty("updateRowIdColumns") Optional<List<DeltaLakeColumnHandle>> updateRowIdColumns,
             @JsonProperty("analyzeHandle") Optional<AnalyzeHandle> analyzeHandle,
@@ -107,7 +106,7 @@ public DeltaLakeTableHandle(
             TupleDomain<DeltaLakeColumnHandle> enforcedPartitionConstraint,
             TupleDomain<DeltaLakeColumnHandle> nonPartitionConstraint,
             Optional<WriteType> writeType,
-            Optional<Set<ColumnHandle>> projectedColumns,
+            Optional<Set<DeltaLakeColumnHandle>> projectedColumns,
             Optional<List<DeltaLakeColumnHandle>> updatedColumns,
             Optional<List<DeltaLakeColumnHandle>> updateRowIdColumns,
             Optional<AnalyzeHandle> analyzeHandle,
@@ -134,7 +133,7 @@ public DeltaLakeTableHandle(
         this.readVersion = readVersion;
     }
 
-    public DeltaLakeTableHandle withProjectedColumns(Set<ColumnHandle> projectedColumns)
+    public DeltaLakeTableHandle withProjectedColumns(Set<DeltaLakeColumnHandle> projectedColumns)
     {
         return new DeltaLakeTableHandle(
                 schemaName,
@@ -245,7 +244,7 @@ public Optional<WriteType> getWriteType()
 
     // Projected columns are not needed on workers
     @JsonIgnore
-    public Optional<Set<ColumnHandle>> getProjectedColumns()
+    public Optional<Set<DeltaLakeColumnHandle>> getProjectedColumns()
     {
         return projectedColumns;
     }

diff --git a/plugin/trino-delta-lake/src/test/java/io/trino/plugin/deltalake/TestDeltaLakeMetadata.java b/plugin/trino-delta-lake/src/test/java/io/trino/plugin/deltalake/TestDeltaLakeMetadata.java
@@ -438,11 +438,11 @@ public Object[][] testApplyProjectionProvider()
 
     @Test(dataProvider = "testApplyProjectionProvider")
     public void testApplyProjection(
-            Set<ColumnHandle> inputProjectedColumns,
+            Set<DeltaLakeColumnHandle> inputProjectedColumns,
             Map<String, ColumnHandle> inputAssignments,
             List<ConnectorExpression> inputProjections,
             List<ConnectorExpression> expectedProjections,
-            Set<ColumnHandle> expectedProjectedColumns,
+            Set<DeltaLakeColumnHandle> expectedProjectedColumns,
             Map<String, ColumnHandle> expectedAssignments)
     {
         DeltaLakeMetadata deltaLakeMetadata = deltaLakeMetadataFactory.create(SESSION.getIdentity());
@@ -455,8 +455,7 @@ public void testApplyProjection(
                         inputAssignments)
                 .get();
 
-        assertThat(((DeltaLakeTableHandle) projection.getHandle())
-                .getProjectedColumns())
+        assertThat(((DeltaLakeTableHandle) projection.getHandle()).getProjectedColumns())
                 .isEqualTo(Optional.of(expectedProjectedColumns));
 
         assertThat(projection.getProjections())
@@ -519,7 +518,7 @@ public void testGetInputInfoForUnPartitionedTable()
         assertThat(deltaLakeMetadata.getInfo(tableHandle)).isEqualTo(Optional.of(new DeltaLakeInputInfo(false)));
     }
 
-    private static DeltaLakeTableHandle createDeltaLakeTableHandle(Set<ColumnHandle> projectedColumns, Set<DeltaLakeColumnHandle> constrainedColumns)
+    private static DeltaLakeTableHandle createDeltaLakeTableHandle(Set<DeltaLakeColumnHandle> projectedColumns, Set<DeltaLakeColumnHandle> constrainedColumns)
     {
         return new DeltaLakeTableHandle(
                 "test_schema_name",