Skip to content

Support adding a new column with id and name column mapping mode in Delta Lake #17236

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 2 commits into from
May 29, 2023
Merged

Support adding a new column with id and name column mapping mode in Delta Lake #17236

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
ce4f47f
Extract a method to get table property value on Delta Lake
ebyhr Apr 29, 2023
1476ca2
Support adding new column with Delta column mapping mode
ebyhr May 18, 2023
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
124 changes: 93 additions & 31 deletions plugin/trino-delta-lake/src/main/java/io/trino/plugin/deltalake/DeltaLakeMetadata.java
View file
Open in desktop

Large diffs are not rendered by default.

9 changes: 9 additions & 0 deletions .../trino-delta-lake/src/main/java/io/trino/plugin/deltalake/DeltaLakeOutputTableHandle.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -38,6 +38,7 @@ public class DeltaLakeOutputTableHandle
private final boolean external;
private final Optional<String> comment;
private final Optional<Boolean> changeDataFeedEnabled;
private final String schemaString;
private final ProtocolEntry protocolEntry;

@JsonCreator
Expand All @@ -50,6 +51,7 @@ public DeltaLakeOutputTableHandle(
@JsonProperty("external") boolean external,
@JsonProperty("comment") Optional<String> comment,
@JsonProperty("changeDataFeedEnabled") Optional<Boolean> changeDataFeedEnabled,
@JsonProperty("schemaString") String schemaString,
@JsonProperty("protocolEntry") ProtocolEntry protocolEntry)
{
this.schemaName = requireNonNull(schemaName, "schemaName is null");
Expand All @@ -60,6 +62,7 @@ public DeltaLakeOutputTableHandle(
this.external = external;
this.comment = requireNonNull(comment, "comment is null");
this.changeDataFeedEnabled = requireNonNull(changeDataFeedEnabled, "changeDataFeedEnabled is null");
this.schemaString = requireNonNull(schemaString, "schemaString is null");
this.protocolEntry = requireNonNull(protocolEntry, "protocolEntry is null");
}

Expand Down Expand Up @@ -120,6 +123,12 @@ public Optional<Boolean> getChangeDataFeedEnabled()
return changeDataFeedEnabled;
}

@JsonProperty
public String getSchemaString()
{
return schemaString;
}

@JsonProperty
public ProtocolEntry getProtocolEntry()
{
Expand Down
9 changes: 1 addition & 8 deletions ...n/trino-delta-lake/src/main/java/io/trino/plugin/deltalake/DeltaLakePageSinkProvider.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -14,7 +14,6 @@
package io.trino.plugin.deltalake;

import com.google.common.collect.ImmutableList;
import com.google.common.collect.ImmutableMap;
import io.airlift.json.JsonCodec;
import io.trino.filesystem.Location;
import io.trino.filesystem.TrinoFileSystemFactory;
Expand Down Expand Up @@ -44,7 +43,6 @@
import java.util.Set;

import static com.google.common.collect.ImmutableList.toImmutableList;
import static com.google.common.collect.ImmutableMap.toImmutableMap;
import static com.google.common.collect.ImmutableSet.toImmutableSet;
import static io.trino.plugin.deltalake.DeltaLakeCdfPageSink.CHANGE_DATA_FOLDER_NAME;
import static io.trino.plugin.deltalake.DeltaLakeCdfPageSink.CHANGE_TYPE_COLUMN_NAME;
Expand All @@ -54,7 +52,6 @@
import static io.trino.plugin.deltalake.transactionlog.DeltaLakeSchemaSupport.ColumnMappingMode.NONE;
import static io.trino.plugin.deltalake.transactionlog.DeltaLakeSchemaSupport.changeDataFeedEnabled;
import static io.trino.plugin.deltalake.transactionlog.DeltaLakeSchemaSupport.extractSchema;
import static io.trino.plugin.deltalake.transactionlog.DeltaLakeSchemaSupport.serializeSchemaAsJson;
import static io.trino.spi.type.VarcharType.VARCHAR;
import static java.util.Objects.requireNonNull;

Expand Down Expand Up @@ -100,11 +97,7 @@ public ConnectorPageSink createPageSink(ConnectorTransactionHandle transactionHa
{
DeltaLakeOutputTableHandle tableHandle = (DeltaLakeOutputTableHandle) outputTableHandle;
DeltaLakeParquetSchemaMapping parquetSchemaMapping = createParquetSchemaMapping(
serializeSchemaAsJson(
tableHandle.getInputColumns(),
ImmutableMap.of(),
tableHandle.getInputColumns().stream().collect(toImmutableMap(DeltaLakeColumnHandle::getColumnName, ignored -> true)),
tableHandle.getInputColumns().stream().collect(toImmutableMap(DeltaLakeColumnHandle::getColumnName, ignored -> ImmutableMap.of()))),
tableHandle.getSchemaString(),
typeManager,
NONE, // TODO support ID and NAME column mapping
tableHandle.getPartitionedBy());
Expand Down
90 changes: 64 additions & 26 deletions ...a-lake/src/main/java/io/trino/plugin/deltalake/transactionlog/DeltaLakeSchemaSupport.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -52,6 +52,8 @@
import java.util.Optional;
import java.util.OptionalInt;
import java.util.Set;
import java.util.UUID;
import java.util.concurrent.atomic.AtomicInteger;
import java.util.function.Function;

import static com.google.common.base.Strings.isNullOrEmpty;
Expand All @@ -76,13 +78,15 @@
import static java.lang.Boolean.parseBoolean;
import static java.lang.String.format;
import static java.util.Locale.ENGLISH;
import static java.util.Objects.requireNonNull;

public final class DeltaLakeSchemaSupport
{
private DeltaLakeSchemaSupport() {}

public static final String APPEND_ONLY_CONFIGURATION_KEY = "delta.appendOnly";
public static final String COLUMN_MAPPING_MODE_CONFIGURATION_KEY = "delta.columnMapping.mode";
public static final String MAX_COLUMN_ID_CONFIGURATION_KEY = "delta.columnMapping.maxColumnId";

// https://github.com/delta-io/delta/blob/master/PROTOCOL.md#valid-feature-names-in-table-features
// TODO: Add support for 'deletionVectors' and 'timestampNTZ' reader features
Expand Down Expand Up @@ -125,6 +129,13 @@ public static ColumnMappingMode getColumnMappingMode(MetadataEntry metadata)
return Enums.getIfPresent(ColumnMappingMode.class, columnMappingMode.toUpperCase(ENGLISH)).or(ColumnMappingMode.UNKNOWN);
}

public static int getMaxColumnId(MetadataEntry metadata)
{
String maxColumnId = metadata.getConfiguration().get(MAX_COLUMN_ID_CONFIGURATION_KEY);
requireNonNull(maxColumnId, () -> MAX_COLUMN_ID_CONFIGURATION_KEY + " metadata configuration property not found");
return Integer.parseInt(maxColumnId);
}

public static List<DeltaLakeColumnHandle> extractPartitionColumns(MetadataEntry metadataEntry, TypeManager typeManager)
{
return extractPartitionColumns(extractSchema(metadataEntry, typeManager), metadataEntry.getCanonicalPartitionColumns());
Expand All @@ -142,21 +153,23 @@ public static List<DeltaLakeColumnHandle> extractPartitionColumns(List<DeltaLake
}

public static String serializeSchemaAsJson(
List<DeltaLakeColumnHandle> columns,
List<String> columnNames,
Map<String, Object> columnTypes,
Map<String, String> columnComments,
Map<String, Boolean> columnNullability,
Map<String, Map<String, Object>> columnMetadata)
{
try {
return OBJECT_MAPPER.writeValueAsString(serializeStructType(columns, columnComments, columnNullability, columnMetadata));
return OBJECT_MAPPER.writeValueAsString(serializeStructType(columnNames, columnTypes, columnComments, columnNullability, columnMetadata));
}
catch (JsonProcessingException e) {
throw new TrinoException(DELTA_LAKE_INVALID_SCHEMA, getLocation(e), "Failed to encode Delta Lake schema", e);
}
}

private static Map<String, Object> serializeStructType(
List<DeltaLakeColumnHandle> columns,
List<String> columnNames,
Map<String, Object> columnTypes,
Map<String, String> columnComments,
Map<String, Boolean> columnNullability,
Map<String, Map<String, Object>> columnMetadata)
Expand All @@ -165,28 +178,25 @@ private static Map<String, Object> serializeStructType(
ImmutableMap.Builder<String, Object> schema = ImmutableMap.builder();

schema.put("type", "struct");
schema.put("fields", columns.stream()
.map(column -> {
String columnName = column.getColumnName();
return serializeStructField(
columnName,
column.getBaseType(),
columnComments.get(columnName),
columnNullability.get(columnName),
columnMetadata.get(columnName));
})
schema.put("fields", columnNames.stream()
.map(columnName -> serializeStructField(
columnName,
columnTypes.get(columnName),
columnComments.get(columnName),
columnNullability.get(columnName),
columnMetadata.get(columnName)))
.collect(toImmutableList()));

return schema.buildOrThrow();
}

private static Map<String, Object> serializeStructField(String name, Type type, @Nullable String comment, @Nullable Boolean nullable, @Nullable Map<String, Object> metadata)
private static Map<String, Object> serializeStructField(String name, Object type, @Nullable String comment, @Nullable Boolean nullable, @Nullable Map<String, Object> metadata)
{
// https://github.com/delta-io/delta/blob/master/PROTOCOL.md#struct-field
ImmutableMap.Builder<String, Object> fieldContents = ImmutableMap.builder();

fieldContents.put("name", name);
fieldContents.put("type", serializeColumnType(type));
fieldContents.put("type", type);
fieldContents.put("nullable", nullable != null ? nullable : true);

ImmutableMap.Builder<String, Object> columnMetadata = ImmutableMap.builder();
Expand All @@ -203,56 +213,79 @@ private static Map<String, Object> serializeStructField(String name, Type type,
return fieldContents.buildOrThrow();
}

private static Object serializeColumnType(Type columnType)
public static Object serializeColumnType(ColumnMappingMode columnMappingMode, AtomicInteger maxColumnId, Type columnType)
{
if (columnType instanceof ArrayType) {
return serializeArrayType((ArrayType) columnType);
return serializeArrayType(columnMappingMode, maxColumnId, (ArrayType) columnType);
}
if (columnType instanceof RowType) {
return serializeStructType((RowType) columnType);
return serializeStructType(columnMappingMode, maxColumnId, (RowType) columnType);
}
if (columnType instanceof MapType) {
return serializeMapType((MapType) columnType);
return serializeMapType(columnMappingMode, maxColumnId, (MapType) columnType);
}
return serializePrimitiveType(columnType);
}

private static Map<String, Object> serializeArrayType(ArrayType arrayType)
private static Map<String, Object> serializeArrayType(ColumnMappingMode columnMappingMode, AtomicInteger maxColumnId, ArrayType arrayType)
{
// https://github.com/delta-io/delta/blob/master/PROTOCOL.md#array-type
ImmutableMap.Builder<String, Object> fields = ImmutableMap.builder();

fields.put("type", "array");
fields.put("elementType", serializeColumnType(arrayType.getElementType()));
fields.put("elementType", serializeColumnType(columnMappingMode, maxColumnId, arrayType.getElementType()));
fields.put("containsNull", true);

return fields.buildOrThrow();
}

private static Map<String, Object> serializeMapType(MapType mapType)
private static Map<String, Object> serializeMapType(ColumnMappingMode columnMappingMode, AtomicInteger maxColumnId, MapType mapType)
{
// https://github.com/delta-io/delta/blob/master/PROTOCOL.md#map-type
ImmutableMap.Builder<String, Object> fields = ImmutableMap.builder();

fields.put("type", "map");
fields.put("keyType", serializeColumnType(mapType.getKeyType()));
fields.put("valueType", serializeColumnType(mapType.getValueType()));
fields.put("keyType", serializeColumnType(columnMappingMode, maxColumnId, mapType.getKeyType()));
fields.put("valueType", serializeColumnType(columnMappingMode, maxColumnId, mapType.getValueType()));
fields.put("valueContainsNull", true);

return fields.buildOrThrow();
}

private static Map<String, Object> serializeStructType(RowType rowType)
private static Map<String, Object> serializeStructType(ColumnMappingMode columnMappingMode, AtomicInteger maxColumnId, RowType rowType)
{
ImmutableMap.Builder<String, Object> fields = ImmutableMap.builder();

fields.put("type", "struct");
fields.put("fields", rowType.getFields().stream()
.map(field -> serializeStructField(field.getName().orElse(null), field.getType(), null, null, null)).collect(toImmutableList()));
.map(field -> {
Object fieldType = serializeColumnType(columnMappingMode, maxColumnId, field.getType());
Map<String, Object> metadata = generateColumnMetadata(columnMappingMode, maxColumnId);
return serializeStructField(field.getName().orElse(null), fieldType, null, null, metadata);
})
.collect(toImmutableList()));

return fields.buildOrThrow();
}

public static Map<String, Object> generateColumnMetadata(ColumnMappingMode columnMappingMode, AtomicInteger maxColumnId)
{
return switch (columnMappingMode) {
case NONE -> {
verify(maxColumnId.get() == 0, "maxColumnId must be 0 for column mapping mode 'none'");
yield ImmutableMap.of();
}
case ID, NAME -> ImmutableMap.<String, Object>builder()
// Set both 'id' and 'physicalName' regardless of the mode https://github.com/delta-io/delta/blob/master/PROTOCOL.md#column-mapping
// > There are two modes of column mapping, by name and by id.
// > In both modes, every column - nested or leaf - is assigned a unique physical name, and a unique 32-bit integer as an id.
.put("delta.columnMapping.id", maxColumnId.incrementAndGet())
.put("delta.columnMapping.physicalName", "col-" + UUID.randomUUID()) // This logic is same as DeltaColumnMapping.generatePhysicalName in Delta Lake
.buildOrThrow();
default -> throw new IllegalArgumentException("Unexpected column mapping mode: " + columnMappingMode);
};
}

private static String serializePrimitiveType(Type type)
{
return serializeSupportedPrimitiveType(type)
Expand Down Expand Up @@ -391,6 +424,11 @@ private static DeltaLakeColumnMetadata mapColumn(TypeManager typeManager, JsonNo
return new DeltaLakeColumnMetadata(columnMetadata, fieldId, physicalName, physicalColumnType);
}

public static Map<String, Object> getColumnTypes(MetadataEntry metadataEntry)
{
return getColumnProperties(metadataEntry, node -> OBJECT_MAPPER.convertValue(node.get("type"), new TypeReference<>(){}));
}

public static Map<String, String> getColumnComments(MetadataEntry metadataEntry)
{
return getColumnProperties(metadataEntry, DeltaLakeSchemaSupport::getComment);
Expand Down
Loading