-
Notifications
You must be signed in to change notification settings - Fork 1.7k
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
[Kernel] Add
partition_value and element_at expressions
- Loading branch information
1 parent
87f80ce
commit 503cde6
Showing
8 changed files
with
581 additions
and
18 deletions.
There are no files selected for viewing
81 changes: 81 additions & 0 deletions
81
kernel/kernel-api/src/main/java/io/delta/kernel/expressions/PartitionValueExpression.java
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,81 @@ | ||
| /* | ||
| * Copyright (2023) The Delta Lake Project Authors. | ||
| * | ||
| * 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 io.delta.kernel.expressions; | ||
|
|
||
| import java.util.Collections; | ||
| import java.util.List; | ||
| import static java.lang.String.format; | ||
| import static java.util.Objects.requireNonNull; | ||
|
|
||
| import io.delta.kernel.annotation.Evolving; | ||
| import io.delta.kernel.types.DataType; | ||
|
|
||
| /** | ||
| * Expression to decode the serialized partition value into partition type value according the | ||
| * <a href=https://github.com/delta-io/delta/blob/master/PROTOCOL.md#partition-value-serialization> | ||
| * Delta Protocol spec</a>. | ||
| * <p> | ||
| * <ul> | ||
| * <li>Name: <code>partition_value</code> | ||
| * <li>Semantic: <code>partition_value(string, datatype)</code>. Decode the partition | ||
| * value of type <i>datatype</i> from the serialized string format.</li> | ||
| * </ul> | ||
| * | ||
| * @since 3.0.0 | ||
| */ | ||
| @Evolving | ||
| public class PartitionValueExpression implements Expression { | ||
| private final DataType partitionValueType; | ||
| private final Expression serializedPartitionValue; | ||
|
|
||
| /** | ||
| * Create {@code partition_value} expression. | ||
| * | ||
| * @param serializedPartitionValue Input expression providing the partition values in | ||
| * serialized format. | ||
| * @param partitionDataType Partition data type to which string partition value is | ||
| * deserialized as according to the Delta Protocol. | ||
| */ | ||
| public PartitionValueExpression( | ||
| Expression serializedPartitionValue, DataType partitionDataType) { | ||
| this.serializedPartitionValue = requireNonNull(serializedPartitionValue); | ||
| this.partitionValueType = requireNonNull(partitionDataType); | ||
| } | ||
|
|
||
| /** | ||
| * Get the expression reference to the serialized partition value. | ||
| */ | ||
| public Expression getInput() { | ||
| return serializedPartitionValue; | ||
| } | ||
|
|
||
| /** | ||
| * Get the data type of the partition value. | ||
| */ | ||
| public DataType getDataType() { | ||
| return partitionValueType; | ||
| } | ||
|
|
||
| @Override | ||
| public List<Expression> getChildren() { | ||
| return Collections.singletonList(serializedPartitionValue); | ||
| } | ||
|
|
||
| @Override | ||
| public String toString() { | ||
| return format("partition_value(%s, %d)", serializedPartitionValue, partitionValueType); | ||
| } | ||
| } |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
133 changes: 133 additions & 0 deletions
133
...aults/src/main/java/io/delta/kernel/defaults/internal/expressions/ElementAtEvaluator.java
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,133 @@ | ||
| /* | ||
| * Copyright (2023) The Delta Lake Project Authors. | ||
| * | ||
| * 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 io.delta.kernel.defaults.internal.expressions; | ||
|
|
||
| import java.util.Arrays; | ||
| import static java.lang.String.format; | ||
|
|
||
| import io.delta.kernel.data.ColumnVector; | ||
| import io.delta.kernel.expressions.Expression; | ||
| import io.delta.kernel.expressions.ScalarExpression; | ||
| import io.delta.kernel.types.DataType; | ||
| import io.delta.kernel.types.MapType; | ||
| import io.delta.kernel.types.StringType; | ||
| import io.delta.kernel.utils.Utils; | ||
|
|
||
| import static io.delta.kernel.defaults.internal.DefaultKernelUtils.checkArgument; | ||
| import static io.delta.kernel.defaults.internal.expressions.ImplicitCastExpression.canCastTo; | ||
|
|
||
| /** | ||
| * Utility methods to evaluate {@code element_at} expression. | ||
| */ | ||
| class ElementAtEvaluator { | ||
| private ElementAtEvaluator() {} | ||
|
|
||
| /** | ||
| * Validate and transform the {@code element_at} expression with given validated and | ||
| * transformed inputs. | ||
| */ | ||
| static ScalarExpression validateAndTransform( | ||
| ScalarExpression elementAt, | ||
| Expression mapInput, | ||
| DataType mapInputType, | ||
| Expression lookupKey, | ||
| DataType lookupKeyType) { | ||
|
|
||
| MapType asMapType = validateSupportedMapType(elementAt, mapInputType); | ||
| DataType keyTypeFromMapInput = asMapType.getKeyType(); | ||
|
|
||
| if (!keyTypeFromMapInput.equivalent(lookupKeyType)) { | ||
| if (canCastTo(lookupKeyType, keyTypeFromMapInput)) { | ||
| lookupKey = new ImplicitCastExpression(lookupKey, keyTypeFromMapInput); | ||
| } else { | ||
| throw new UnsupportedOperationException(format( | ||
| "%s: lookup key type (%s) is different from the map key type (%s)", | ||
| elementAt, lookupKeyType, asMapType.getKeyType())); | ||
| } | ||
| } | ||
| return new ScalarExpression(elementAt.getName(), Arrays.asList(mapInput, lookupKey)); | ||
| } | ||
|
|
||
| /** | ||
| * Utility method to evaluate the {@code element_at} on given map and key vectors. | ||
| * @param map {@link ColumnVector} of {@code map(string, string)} type. | ||
| * @param lookupKey {@link ColumnVector} of {@code string} type. | ||
| * @return | ||
| */ | ||
| static ColumnVector eval(ColumnVector map, ColumnVector lookupKey) { | ||
| return new ColumnVector() { | ||
| // Store the last lookup value to avoid multiple looks up for same row id. | ||
| private int lastLookupRowId = -1; | ||
| private Object lastLookupValue = null; | ||
|
|
||
| @Override | ||
| public DataType getDataType() { | ||
| return ((MapType) map.getDataType()).getValueType(); | ||
| } | ||
|
|
||
| @Override | ||
| public int getSize() { | ||
| return map.getSize(); | ||
| } | ||
|
|
||
| @Override | ||
| public void close() { | ||
| Utils.closeCloseables(map, lookupKey); | ||
| } | ||
|
|
||
| @Override | ||
| public boolean isNullAt(int rowId) { | ||
| if (rowId == lastLookupRowId) { | ||
| return lastLookupValue == null; | ||
| } | ||
| return map.isNullAt(rowId) || lookupValue(rowId) == null; | ||
| } | ||
|
|
||
| @Override | ||
| public String getString(int rowId) { | ||
| lookupValue(rowId); | ||
| return lastLookupValue == null ? null : (String) lastLookupValue; | ||
| } | ||
|
|
||
| private Object lookupValue(int rowId) { | ||
| if (rowId == lastLookupRowId) { | ||
| return lastLookupValue; | ||
| } | ||
| // TODO: this needs to be updated after the new way of accessing the complex | ||
| // types is merged. | ||
| lastLookupRowId = rowId; | ||
| String keyValue = lookupKey.getString(rowId); | ||
| lastLookupValue = map.getMap(rowId).get(keyValue); | ||
| return lastLookupValue; | ||
| } | ||
| }; | ||
| } | ||
|
|
||
| private static MapType validateSupportedMapType(Expression elementAt, DataType mapInputType) { | ||
| checkArgument( | ||
| mapInputType instanceof MapType, | ||
| "expected a map type input as first argument: " + elementAt); | ||
| MapType asMapType = (MapType) mapInputType; | ||
| // TODO: we may extend type support in future, but for the need is just a look | ||
| // in map(string, string). | ||
| if (asMapType.getKeyType().equivalent(StringType.INSTANCE) && | ||
| asMapType.getValueType().equivalent(StringType.INSTANCE)) { | ||
| return asMapType; | ||
| } | ||
| throw new UnsupportedOperationException( | ||
| format("%s: Supported only on type map(string, string) input data", elementAt)); | ||
| } | ||
| } |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Oops, something went wrong.