diff --git a/paimon-common/src/main/java/org/apache/paimon/data/shredding/MapSelectedKeysMetadataUtils.java b/paimon-common/src/main/java/org/apache/paimon/data/shredding/MapSelectedKeysMetadataUtils.java new file mode 100644 index 000000000000..6a4163d673a4 --- /dev/null +++ b/paimon-common/src/main/java/org/apache/paimon/data/shredding/MapSelectedKeysMetadataUtils.java @@ -0,0 +1,102 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you 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 org.apache.paimon.data.shredding; + +import org.apache.paimon.types.DataField; +import org.apache.paimon.types.DataType; +import org.apache.paimon.types.RowType; + +import javax.annotation.Nullable; + +import java.util.ArrayList; +import java.util.LinkedHashSet; +import java.util.List; +import java.util.Set; + +import static org.apache.paimon.utils.Preconditions.checkArgument; +import static org.apache.paimon.utils.Preconditions.checkNotNull; + +/** + * Utils for marking a temporary read {@link RowType} field as selected-key MAP output. Uses + * description field in DataField to encode selected-key metadata. + * + *

Example: __PAIMON_MAP_SELECTED_KEYS:key1;key2 + */ +public class MapSelectedKeysMetadataUtils { + + public static final String METADATA_KEY = "__PAIMON_MAP_SELECTED_KEYS:"; + public static final String KEY_DELIMITER = ";"; + + private MapSelectedKeysMetadataUtils() {} + + public static String buildMapSelectedKeysMetadata(List keys) { + List normalizedKeys = normalizeKeys(keys); + return METADATA_KEY + String.join(KEY_DELIMITER, normalizedKeys); + } + + public static boolean isMapSelectedKeysMetadata(@Nullable String description) { + return description != null && description.startsWith(METADATA_KEY); + } + + public static boolean isMapSelectedKeysField(DataField field) { + return isMapSelectedKeysMetadata(field.description()) && field.type() instanceof RowType; + } + + public static DataField withSelectedKeys(DataField field, DataType rowType, List keys) { + checkArgument(rowType instanceof RowType, "Selected-key MAP read type must be ROW."); + return field.newType(rowType).newDescription(buildMapSelectedKeysMetadata(keys)); + } + + public static List selectedKeys(String description) { + checkArgument( + isMapSelectedKeysMetadata(description), + "Invalid selected-key MAP metadata: %s", + description); + String encoded = description.substring(METADATA_KEY.length()); + + String[] parts = encoded.split(KEY_DELIMITER, -1); + List keys = new ArrayList<>(parts.length); + for (String part : parts) { + keys.add(part); + } + return keys; + } + + private static List normalizeKeys(List keys) { + checkNotNull(keys, "Selected keys must not be null."); + checkArgument(!keys.isEmpty(), "Selected keys must not be empty."); + + Set seenKeys = new LinkedHashSet<>(); + for (String key : keys) { + checkNotNull(key, "Selected key must not be null."); + checkArgument( + !key.contains(KEY_DELIMITER), + "Selected key must not contain '%s': %s", + KEY_DELIMITER, + key); + checkArgument( + !key.startsWith(METADATA_KEY), + "Selected key must not start with metadata prefix: %s", + key); + checkArgument(!seenKeys.contains(key), "Selected key must not be duplicated: %s", key); + seenKeys.add(key); + } + return new ArrayList<>(seenKeys); + } +} diff --git a/paimon-common/src/test/java/org/apache/paimon/data/shredding/MapSelectedKeysMetadataUtilsTest.java b/paimon-common/src/test/java/org/apache/paimon/data/shredding/MapSelectedKeysMetadataUtilsTest.java new file mode 100644 index 000000000000..c921ee75284e --- /dev/null +++ b/paimon-common/src/test/java/org/apache/paimon/data/shredding/MapSelectedKeysMetadataUtilsTest.java @@ -0,0 +1,131 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you 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 org.apache.paimon.data.shredding; + +import org.apache.paimon.types.DataField; +import org.apache.paimon.types.DataTypes; +import org.apache.paimon.types.RowType; + +import org.junit.jupiter.api.Test; + +import java.util.Arrays; +import java.util.Collections; + +import static org.assertj.core.api.Assertions.assertThat; +import static org.assertj.core.api.Assertions.assertThatThrownBy; + +/** Tests for {@link MapSelectedKeysMetadataUtils}. */ +class MapSelectedKeysMetadataUtilsTest { + + @Test + void testBuildAndParseMetadata() { + String metadata = + MapSelectedKeysMetadataUtils.buildMapSelectedKeysMetadata( + Arrays.asList("key1", "key2")); + + assertThat(metadata).isEqualTo("__PAIMON_MAP_SELECTED_KEYS:key1;key2"); + assertThat(MapSelectedKeysMetadataUtils.isMapSelectedKeysMetadata(metadata)).isTrue(); + assertThat(MapSelectedKeysMetadataUtils.selectedKeys(metadata)) + .containsExactly("key1", "key2"); + } + + @Test + void testBuildAndParseMetadataWithEmptyKey() { + String metadata = + MapSelectedKeysMetadataUtils.buildMapSelectedKeysMetadata( + Arrays.asList("key1", "", "key2")); + + String expectedMetadata = String.join(";", "__PAIMON_MAP_SELECTED_KEYS:key1", "", "key2"); + assertThat(metadata).isEqualTo(expectedMetadata); + assertThat(MapSelectedKeysMetadataUtils.selectedKeys(metadata)) + .containsExactly("key1", "", "key2"); + assertThat(MapSelectedKeysMetadataUtils.selectedKeys("__PAIMON_MAP_SELECTED_KEYS:")) + .containsExactly(""); + } + + @Test + void testBuildMetadataRejectsDuplicateKeys() { + assertThatThrownBy( + () -> + MapSelectedKeysMetadataUtils.buildMapSelectedKeysMetadata( + Arrays.asList("key1", "key2", "key1"))) + .hasMessageContaining("Selected key must not be duplicated: key1"); + assertThatThrownBy( + () -> + MapSelectedKeysMetadataUtils.buildMapSelectedKeysMetadata( + Arrays.asList("key1", "", ""))) + .hasMessageContaining("Selected key must not be duplicated"); + } + + @Test + void testWithSelectedKeys() { + DataField field = + DataTypes.FIELD( + 1, + "attrs", + DataTypes.MAP(DataTypes.STRING(), DataTypes.BIGINT()), + "user comment"); + RowType selectedType = + DataTypes.ROW( + DataTypes.FIELD(0, "0", DataTypes.BIGINT()), + DataTypes.FIELD(1, "1", DataTypes.BIGINT())); + + DataField rewritten = + MapSelectedKeysMetadataUtils.withSelectedKeys( + field, selectedType, Arrays.asList("key1", "key2")); + + assertThat(rewritten.id()).isEqualTo(1); + assertThat(rewritten.name()).isEqualTo("attrs"); + assertThat(rewritten.type()).isEqualTo(selectedType); + assertThat(rewritten.description()).isEqualTo("__PAIMON_MAP_SELECTED_KEYS:key1;key2"); + assertThat(MapSelectedKeysMetadataUtils.isMapSelectedKeysField(rewritten)).isTrue(); + } + + @Test + void testErrors() { + assertThatThrownBy( + () -> + MapSelectedKeysMetadataUtils.buildMapSelectedKeysMetadata( + Collections.emptyList())) + .hasMessageContaining("Selected keys must not be empty."); + assertThatThrownBy( + () -> + MapSelectedKeysMetadataUtils.buildMapSelectedKeysMetadata( + Arrays.asList("key1", null))) + .hasMessageContaining("Selected key must not be null."); + assertThatThrownBy( + () -> + MapSelectedKeysMetadataUtils.buildMapSelectedKeysMetadata( + Arrays.asList("key;1"))) + .hasMessageContaining("Selected key must not contain ';'"); + assertThatThrownBy(() -> MapSelectedKeysMetadataUtils.selectedKeys("invalid")) + .hasMessageContaining("Invalid selected-key MAP metadata"); + assertThatThrownBy( + () -> + MapSelectedKeysMetadataUtils.withSelectedKeys( + DataTypes.FIELD( + 1, + "attrs", + DataTypes.MAP( + DataTypes.STRING(), DataTypes.BIGINT())), + DataTypes.BIGINT(), + Arrays.asList("key1"))) + .hasMessageContaining("Selected-key MAP read type must be ROW."); + } +} diff --git a/paimon-spark/paimon-spark-3.2/src/main/scala/org/apache/paimon/spark/PaimonScan.scala b/paimon-spark/paimon-spark-3.2/src/main/scala/org/apache/paimon/spark/PaimonScan.scala index e1fd2581867e..fab2cfe8612e 100644 --- a/paimon-spark/paimon-spark-3.2/src/main/scala/org/apache/paimon/spark/PaimonScan.scala +++ b/paimon-spark/paimon-spark-3.2/src/main/scala/org/apache/paimon/spark/PaimonScan.scala @@ -36,5 +36,6 @@ case class PaimonScan( override val pushedHybridSearch: Option[HybridSearch] = None, override val pushedFullTextSearch: Option[FullTextSearch] = None, override val pushedVariantExtractions: Map[Seq[String], Seq[VariantExtractionInfo]] = Map.empty, + override val pushedMapSelectedKeys: Map[String, Seq[String]] = Map.empty, bucketedScanDisabled: Boolean = true) extends PaimonBaseScan(table) {} diff --git a/paimon-spark/paimon-spark-3.2/src/main/scala/org/apache/paimon/spark/catalyst/optimizer/PushDownMapSelectedKeys.scala b/paimon-spark/paimon-spark-3.2/src/main/scala/org/apache/paimon/spark/catalyst/optimizer/PushDownMapSelectedKeys.scala new file mode 100644 index 000000000000..bc4f413e46be --- /dev/null +++ b/paimon-spark/paimon-spark-3.2/src/main/scala/org/apache/paimon/spark/catalyst/optimizer/PushDownMapSelectedKeys.scala @@ -0,0 +1,34 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you 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 org.apache.paimon.spark.catalyst.optimizer + +import org.apache.paimon.spark.PaimonScan + +import org.apache.spark.sql.catalyst.expressions.AttributeReference +import org.apache.spark.sql.execution.datasources.v2.DataSourceV2ScanRelation + +object PushDownMapSelectedKeys extends PushDownMapSelectedKeysBase { + + override protected def copyDataSourceV2ScanRelation( + relation: DataSourceV2ScanRelation, + scan: PaimonScan, + output: Seq[AttributeReference]): DataSourceV2ScanRelation = { + DataSourceV2ScanRelation(relation.relation, scan, output) + } +} diff --git a/paimon-spark/paimon-spark-3.3/src/main/scala/org/apache/paimon/spark/PaimonScan.scala b/paimon-spark/paimon-spark-3.3/src/main/scala/org/apache/paimon/spark/PaimonScan.scala index 6ecc3dd847b7..9cc7e409c7fe 100644 --- a/paimon-spark/paimon-spark-3.3/src/main/scala/org/apache/paimon/spark/PaimonScan.scala +++ b/paimon-spark/paimon-spark-3.3/src/main/scala/org/apache/paimon/spark/PaimonScan.scala @@ -40,6 +40,7 @@ case class PaimonScan( override val pushedHybridSearch: Option[HybridSearch] = None, override val pushedFullTextSearch: Option[FullTextSearch] = None, override val pushedVariantExtractions: Map[Seq[String], Seq[VariantExtractionInfo]] = Map.empty, + override val pushedMapSelectedKeys: Map[String, Seq[String]] = Map.empty, bucketedScanDisabled: Boolean = false) extends PaimonBaseScan(table) with SupportsReportPartitioning { diff --git a/paimon-spark/paimon-spark-3.3/src/main/scala/org/apache/paimon/spark/catalyst/optimizer/PushDownMapSelectedKeys.scala b/paimon-spark/paimon-spark-3.3/src/main/scala/org/apache/paimon/spark/catalyst/optimizer/PushDownMapSelectedKeys.scala new file mode 100644 index 000000000000..823c32a30204 --- /dev/null +++ b/paimon-spark/paimon-spark-3.3/src/main/scala/org/apache/paimon/spark/catalyst/optimizer/PushDownMapSelectedKeys.scala @@ -0,0 +1,34 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you 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 org.apache.paimon.spark.catalyst.optimizer + +import org.apache.paimon.spark.PaimonScan + +import org.apache.spark.sql.catalyst.expressions.AttributeReference +import org.apache.spark.sql.execution.datasources.v2.DataSourceV2ScanRelation + +object PushDownMapSelectedKeys extends PushDownMapSelectedKeysBase { + + override protected def copyDataSourceV2ScanRelation( + relation: DataSourceV2ScanRelation, + scan: PaimonScan, + output: Seq[AttributeReference]): DataSourceV2ScanRelation = { + DataSourceV2ScanRelation(relation.relation, scan, output, relation.keyGroupedPartitioning) + } +} diff --git a/paimon-spark/paimon-spark-3.4/src/main/scala/org/apache/paimon/spark/catalyst/optimizer/PushDownMapSelectedKeys.scala b/paimon-spark/paimon-spark-3.4/src/main/scala/org/apache/paimon/spark/catalyst/optimizer/PushDownMapSelectedKeys.scala new file mode 100644 index 000000000000..88e6347cee2e --- /dev/null +++ b/paimon-spark/paimon-spark-3.4/src/main/scala/org/apache/paimon/spark/catalyst/optimizer/PushDownMapSelectedKeys.scala @@ -0,0 +1,39 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you 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 org.apache.paimon.spark.catalyst.optimizer + +import org.apache.paimon.spark.PaimonScan + +import org.apache.spark.sql.catalyst.expressions.AttributeReference +import org.apache.spark.sql.execution.datasources.v2.DataSourceV2ScanRelation + +object PushDownMapSelectedKeys extends PushDownMapSelectedKeysBase { + + override protected def copyDataSourceV2ScanRelation( + relation: DataSourceV2ScanRelation, + scan: PaimonScan, + output: Seq[AttributeReference]): DataSourceV2ScanRelation = { + DataSourceV2ScanRelation( + relation.relation, + scan, + output, + relation.keyGroupedPartitioning, + relation.ordering) + } +} diff --git a/paimon-spark/paimon-spark-3.5/src/main/scala/org/apache/paimon/spark/catalyst/optimizer/PushDownMapSelectedKeys.scala b/paimon-spark/paimon-spark-3.5/src/main/scala/org/apache/paimon/spark/catalyst/optimizer/PushDownMapSelectedKeys.scala new file mode 100644 index 000000000000..88e6347cee2e --- /dev/null +++ b/paimon-spark/paimon-spark-3.5/src/main/scala/org/apache/paimon/spark/catalyst/optimizer/PushDownMapSelectedKeys.scala @@ -0,0 +1,39 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you 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 org.apache.paimon.spark.catalyst.optimizer + +import org.apache.paimon.spark.PaimonScan + +import org.apache.spark.sql.catalyst.expressions.AttributeReference +import org.apache.spark.sql.execution.datasources.v2.DataSourceV2ScanRelation + +object PushDownMapSelectedKeys extends PushDownMapSelectedKeysBase { + + override protected def copyDataSourceV2ScanRelation( + relation: DataSourceV2ScanRelation, + scan: PaimonScan, + output: Seq[AttributeReference]): DataSourceV2ScanRelation = { + DataSourceV2ScanRelation( + relation.relation, + scan, + output, + relation.keyGroupedPartitioning, + relation.ordering) + } +} diff --git a/paimon-spark/paimon-spark-common/src/main/scala/org/apache/paimon/spark/PaimonScan.scala b/paimon-spark/paimon-spark-common/src/main/scala/org/apache/paimon/spark/PaimonScan.scala index 412d92b7f484..1f0b1b75b038 100644 --- a/paimon-spark/paimon-spark-common/src/main/scala/org/apache/paimon/spark/PaimonScan.scala +++ b/paimon-spark/paimon-spark-common/src/main/scala/org/apache/paimon/spark/PaimonScan.scala @@ -46,6 +46,7 @@ case class PaimonScan( override val pushedHybridSearch: Option[HybridSearch] = None, override val pushedFullTextSearch: Option[FullTextSearch] = None, override val pushedVariantExtractions: Map[Seq[String], Seq[VariantExtractionInfo]] = Map.empty, + override val pushedMapSelectedKeys: Map[String, Seq[String]] = Map.empty, bucketedScanDisabled: Boolean = false) extends PaimonBaseScan(table) with SupportsReportPartitioning diff --git a/paimon-spark/paimon-spark-common/src/main/scala/org/apache/paimon/spark/catalyst/optimizer/PushDownMapSelectedKeys.scala b/paimon-spark/paimon-spark-common/src/main/scala/org/apache/paimon/spark/catalyst/optimizer/PushDownMapSelectedKeys.scala new file mode 100644 index 000000000000..16db25531fb5 --- /dev/null +++ b/paimon-spark/paimon-spark-common/src/main/scala/org/apache/paimon/spark/catalyst/optimizer/PushDownMapSelectedKeys.scala @@ -0,0 +1,289 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you 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 org.apache.paimon.spark.catalyst.optimizer + +import org.apache.paimon.CoreOptions +import org.apache.paimon.CoreOptions.MapStorageLayout +import org.apache.paimon.data.shredding.MapSelectedKeysMetadataUtils +import org.apache.paimon.data.shredding.MapSharedShreddingUtils +import org.apache.paimon.spark.PaimonScan + +import org.apache.spark.sql.catalyst.expressions.{Alias, Attribute, AttributeReference, Expression, ExprId, GetMapValue, GetStructField, Literal, NamedExpression} +import org.apache.spark.sql.catalyst.plans.logical.{LogicalPlan, Project} +import org.apache.spark.sql.catalyst.rules.Rule +import org.apache.spark.sql.execution.datasources.v2.DataSourceV2ScanRelation +import org.apache.spark.sql.types.{MapType, StringType, StructField, StructType} + +import scala.collection.mutable + +/** + * Pushes literal-key MAP access on shared-shredding MAP columns into Paimon scan read type. + * + *

Example: {@code SELECT id, attrs['key1'], attrs['key2'] FROM T} becomes a scan returning + * {@code attrs} as {@code struct<0:valueType,1:valueType>} and a project reading struct fields. + */ +object PushDownMapSelectedKeys extends Rule[LogicalPlan] { + + override def apply(plan: LogicalPlan): LogicalPlan = { + plan + } +} + +abstract class PushDownMapSelectedKeysBase extends Rule[LogicalPlan] { + + override def apply(plan: LogicalPlan): LogicalPlan = plan.transform { + case project @ Project(projectList, relation: DataSourceV2ScanRelation) => + relation.scan match { + case scan: PaimonScan => + rewriteProject(project, projectList, relation, scan).getOrElse(project) + case _ => project + } + } + + private def rewriteProject( + project: Project, + projectList: Seq[NamedExpression], + relation: DataSourceV2ScanRelation, + scan: PaimonScan): Option[LogicalPlan] = { + val candidates = collectCandidates(projectList, scan) + if (candidates.isEmpty) { + return None + } + + val outputByExprId = relation.output.map(a => a.exprId -> a).toMap + val selectedMaps = mutable.LinkedHashMap.empty[ExprId, SelectedMap] + candidates.foreach { + candidate => + outputByExprId.get(candidate.rootExprId).foreach { + root => + val selectedMap = + selectedMaps.getOrElseUpdate( + root.exprId, + SelectedMap( + root, + candidate.fieldName, + candidate.mapType, + mutable.ArrayBuffer.empty)) + if (!selectedMap.keys.contains(candidate.key)) { + selectedMap.keys.append(candidate.key) + } + } + } + if (selectedMaps.isEmpty) { + return None + } + + val selectedByExprId = selectedMaps.toMap + val attrToRewritten = selectedByExprId.map { + case (exprId, selected) => + exprId -> selected.root.withDataType(selected.structType).asInstanceOf[AttributeReference] + } + + val rewriteState = RewriteState(selectedByExprId, attrToRewritten) + val rewrittenProjectList = projectList.map(rewriteExpression(_, rewriteState)) + if (rewriteState.failed) { + return None + } + + val pushedMapSelectedKeys = scan.pushedMapSelectedKeys ++ selectedMaps.values.map { + selected => selected.fieldName -> selected.keys.toSeq + }.toMap + val rewrittenScan = scan.copy(pushedMapSelectedKeys = pushedMapSelectedKeys) + val rewrittenOutput = + relation.output.map(attr => attrToRewritten.getOrElse(attr.exprId, attr)) + val rewrittenRelation = copyDataSourceV2ScanRelation(relation, rewrittenScan, rewrittenOutput) + Some(project.copy(projectList = rewrittenProjectList, child = rewrittenRelation)) + } + + protected def copyDataSourceV2ScanRelation( + relation: DataSourceV2ScanRelation, + scan: PaimonScan, + output: Seq[AttributeReference]): DataSourceV2ScanRelation + + private def collectCandidates( + projectList: Seq[NamedExpression], + scan: PaimonScan): Seq[MapKeyCandidate] = { + val candidates = mutable.ArrayBuffer.empty[MapKeyCandidate] + projectList.foreach { + expr => + expr.foreach { + case mapKeyValueAccess(access) if canPushDown(scan, access) => + candidates.append( + MapKeyCandidate(access.root.exprId, access.fieldName, access.mapType, access.key)) + case _ => + } + } + candidates.toSeq + } + + private def rewriteExpression( + expression: NamedExpression, + state: RewriteState): NamedExpression = { + expression match { + case alias: Alias => + alias.child match { + case mapKeyValueAccess(access) if state.selectedByExprId.contains(access.root.exprId) => + rewriteMapKeyAccess(access, state) match { + case Some(rewritten) => + alias.withNewChildren(Seq(rewritten)).asInstanceOf[NamedExpression] + case None => expression + } + case _ => + if (hasSelectedMapReference(alias, state)) { + state.failed = true + } + expression + } + case attr: Attribute if state.hasRootSelection(attr.exprId) => + state.failed = true + expression + case _ => + if (hasSelectedMapReference(expression, state)) { + state.failed = true + } + expression + } + } + + private def rewriteMapKeyAccess(access: MapKeyAccess, state: RewriteState): Option[Expression] = { + val selected = state.selectedByExprId(access.root.exprId) + val ordinal = selected.keys.indexOf(access.key) + if (ordinal < 0) { + state.failed = true + None + } else { + Some(GetStructField(state.attrToRewritten(access.root.exprId), ordinal, Some(access.key))) + } + } + + private def hasSelectedMapReference(expression: Expression, state: RewriteState): Boolean = { + var found = false + expression.foreach { + case mapFieldReference(access) if state.selectedByExprId.contains(access.root.exprId) => + found = true + case _ => + } + found + } + + private def canPushDown(scan: PaimonScan, access: MapKeyAccess): Boolean = { + if (!canEncodeKey(access.key)) { + return false + } + + access.mapType match { + case MapType(StringType, _, _) => + fieldType(scan.table.rowType(), access.fieldName) match { + case Some(mapType: org.apache.paimon.types.MapType) if isStringKeyMap(mapType) => + CoreOptions.fromMap(scan.table.options()).mapStorageLayout(access.fieldName) == + MapStorageLayout.SHARED_SHREDDING + case _ => false + } + case _ => false + } + } + + private def canEncodeKey(key: String): Boolean = { + !key.contains(MapSelectedKeysMetadataUtils.KEY_DELIMITER) && + !key.startsWith(MapSelectedKeysMetadataUtils.METADATA_KEY) + } + + private def isStringKeyMap(mapType: org.apache.paimon.types.MapType): Boolean = { + MapSharedShreddingUtils.isShreddingKeyMap(mapType) + } + + private object mapKeyValueAccess { + def unapply(expression: Expression): Option[MapKeyAccess] = { + expression match { + case GetMapValue(mapFieldReference(access), StringLiteral(key)) => + Some(access.copy(key = key)) + case e if e.prettyName == "element_at" && e.children.length == 2 => + (e.children.head, e.children(1)) match { + case (mapFieldReference(access), StringLiteral(key)) => Some(access.copy(key = key)) + case _ => None + } + case _ => None + } + } + } + + private object mapFieldReference { + def unapply(expression: Expression): Option[MapKeyAccess] = { + expression match { + case attr: Attribute => + attr.dataType match { + case mapType: MapType => Some(MapKeyAccess(attr, attr.name, mapType, "")) + case _ => None + } + case _ => None + } + } + } + + private object StringLiteral { + def unapply(expression: Expression): Option[String] = { + expression match { + case Literal(value, StringType) if value != null => Some(value.toString) + case _ => None + } + } + } + + private def fieldType( + rowType: org.apache.paimon.types.RowType, + fieldName: String): Option[org.apache.paimon.types.DataType] = { + if (!rowType.containsField(fieldName)) { + None + } else { + Some(rowType.getField(fieldName).`type`()) + } + } + + private case class MapKeyCandidate( + rootExprId: ExprId, + fieldName: String, + mapType: MapType, + key: String) + + private case class MapKeyAccess(root: Attribute, fieldName: String, mapType: MapType, key: String) + + private case class SelectedMap( + root: Attribute, + fieldName: String, + mapType: MapType, + keys: mutable.ArrayBuffer[String]) { + def structType: StructType = { + val MapType(_, valueType, _) = mapType + StructType(keys.zipWithIndex.map { + case (_, ordinal) => + StructField(ordinal.toString, valueType, nullable = true) + }.toSeq) + } + } + + private case class RewriteState( + selectedByExprId: Map[ExprId, SelectedMap], + attrToRewritten: Map[ExprId, AttributeReference]) { + var failed: Boolean = false + + def hasRootSelection(exprId: ExprId): Boolean = { + selectedByExprId.contains(exprId) + } + } +} diff --git a/paimon-spark/paimon-spark-common/src/main/scala/org/apache/paimon/spark/extensions/PaimonSparkSessionExtensions.scala b/paimon-spark/paimon-spark-common/src/main/scala/org/apache/paimon/spark/extensions/PaimonSparkSessionExtensions.scala index 61481e201c0e..388889bbeaaa 100644 --- a/paimon-spark/paimon-spark-common/src/main/scala/org/apache/paimon/spark/extensions/PaimonSparkSessionExtensions.scala +++ b/paimon-spark/paimon-spark-common/src/main/scala/org/apache/paimon/spark/extensions/PaimonSparkSessionExtensions.scala @@ -101,6 +101,8 @@ class PaimonSparkSessionExtensions extends (SparkSessionExtensions => Unit) { // optimization rules extensions.injectOptimizerRule(spark => ReplacePaimonFunctions(spark)) extensions.injectOptimizerRule(_ => OptimizeMetadataOnlyDeleteFromPaimonTable) + // TODO: Enable MAP selected-key pushdown after core reader supports + // __PAIMON_MAP_SELECTED_KEYS read type. extensions.injectOptimizerRule(_ => MergePaimonScalarSubqueries) extensions.injectOptimizerRule(_ => PushDownLateralVectorSearchFilter) diff --git a/paimon-spark/paimon-spark-common/src/main/scala/org/apache/paimon/spark/read/BaseScan.scala b/paimon-spark/paimon-spark-common/src/main/scala/org/apache/paimon/spark/read/BaseScan.scala index b36a0f1f02f3..088b785c7b2e 100644 --- a/paimon-spark/paimon-spark-common/src/main/scala/org/apache/paimon/spark/read/BaseScan.scala +++ b/paimon-spark/paimon-spark-common/src/main/scala/org/apache/paimon/spark/read/BaseScan.scala @@ -54,6 +54,7 @@ trait BaseScan extends Scan with SupportsReportStatistics with Logging { def pushedHybridSearch: Option[HybridSearch] = None def pushedFullTextSearch: Option[FullTextSearch] = None def pushedVariantExtractions: Map[Seq[String], Seq[VariantExtractionInfo]] = Map.empty + def pushedMapSelectedKeys: Map[String, Seq[String]] = Map.empty // Runtime push down val pushedRuntimePartitionFilters: ListBuffer[PartitionPredicate] = ListBuffer.empty @@ -80,7 +81,10 @@ trait BaseScan extends Scan with SupportsReportStatistics with Logging { } } - /** Pruned read RowType, with variant fields rewritten if variant pushdown was accepted. */ + /** + * Pruned read RowType, with variant fields rewritten if variant pushdown was accepted, with + * accepted nested field pushdowns rewritten. + */ private[paimon] val (readTableRowType, metadataFields) = { requiredSchema.fields.foreach(f => checkMetadataColumn(f.name)) val (_requiredTableFields, _metadataFields) = @@ -90,7 +94,10 @@ trait BaseScan extends Scan with SupportsReportStatistics with Logging { val withVariants = if (pushedVariantExtractions.isEmpty) pruned else VariantPushDownUtils.rewriteRowType(pruned, pushedVariantExtractions) - (withVariants, _metadataFields) + val withMapSelectedKeys = + if (pushedMapSelectedKeys.isEmpty) withVariants + else MapSelectedKeysPushDownUtils.rewriteRowType(withVariants, pushedMapSelectedKeys) + (withMapSelectedKeys, _metadataFields) } private def checkMetadataColumn(fieldName: String): Unit = { @@ -198,6 +205,13 @@ trait BaseScan extends Scan with SupportsReportStatistics with Logging { .describeRewrittenRowType(readTableRowType) .map(s => s", PushedVariants: [$s]") .getOrElse("") + val pushedMapSelectedKeysStr = + if (pushedMapSelectedKeys.isEmpty) "" + else + MapSelectedKeysPushDownUtils + .describeRewrittenRowType(readTableRowType) + .map(s => s", PushedMapSelectedKeys: [$s]") + .getOrElse("") s"${getClass.getSimpleName}: [${table.name}]" + pushedPartitionFiltersStr + pushedRuntimePartitionFiltersStr + @@ -207,6 +221,7 @@ trait BaseScan extends Scan with SupportsReportStatistics with Logging { pushedVectorSearch.map(vs => s", VectorSearch: [$vs]").getOrElse("") + pushedHybridSearch.map(hs => s", HybridSearch: [$hs]").getOrElse("") + pushedFullTextSearch.map(fts => s", FullTextSearch: [$fts]").getOrElse("") + - pushedVariantsStr + pushedVariantsStr + + pushedMapSelectedKeysStr } } diff --git a/paimon-spark/paimon-spark-common/src/main/scala/org/apache/paimon/spark/read/MapSelectedKeysPushDownUtils.scala b/paimon-spark/paimon-spark-common/src/main/scala/org/apache/paimon/spark/read/MapSelectedKeysPushDownUtils.scala new file mode 100644 index 000000000000..bb131d2e78ab --- /dev/null +++ b/paimon-spark/paimon-spark-common/src/main/scala/org/apache/paimon/spark/read/MapSelectedKeysPushDownUtils.scala @@ -0,0 +1,71 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you 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 org.apache.paimon.spark.read + +import org.apache.paimon.data.shredding.MapSelectedKeysMetadataUtils +import org.apache.paimon.types.{DataField, MapType, RowType} + +import scala.collection.JavaConverters._ +import scala.collection.mutable + +/** Shared-shredding MAP selected-key pushdown business logic. */ +object MapSelectedKeysPushDownUtils { + + /** Replace accepted top-level MAP fields with selected-key ROW fields in the read type. */ + def rewriteRowType(rt: RowType, accepted: Map[String, Seq[String]]): RowType = { + val newFields = rt.getFields.asScala.map(f => rewriteField(f, accepted)).asJava + new RowType(rt.isNullable, newFields) + } + + private def rewriteField(field: DataField, accepted: Map[String, Seq[String]]): DataField = { + accepted.get(field.name()) match { + case Some(keys) => + field.`type`() match { + case mapType: MapType => + val selectedFields = keys.zipWithIndex.map { + case (_, ordinal) => + new DataField(ordinal, ordinal.toString, mapType.getValueType.copy(true)) + }.asJava + MapSelectedKeysMetadataUtils.withSelectedKeys( + field, + new RowType(field.`type`().isNullable, selectedFields), + keys.asJava) + case _ => field + } + case None => field + } + } + + /** "col=[key1,key2]" rendering for `Scan.description()`. */ + def describeRewrittenRowType(rt: RowType): Option[String] = { + val parts = mutable.ArrayBuffer.empty[String] + collectSelectedKeys(rt, parts) + if (parts.isEmpty) None else Some(parts.mkString(", ")) + } + + private def collectSelectedKeys(rt: RowType, out: mutable.ArrayBuffer[String]): Unit = { + rt.getFields.asScala.foreach { + field => + if (MapSelectedKeysMetadataUtils.isMapSelectedKeysField(field)) { + val keys = MapSelectedKeysMetadataUtils.selectedKeys(field.description()).asScala + out.append(s"${field.name()}=[${keys.mkString(",")}]") + } + } + } +} diff --git a/paimon-spark/paimon-spark-common/src/test/scala/org/apache/paimon/spark/read/MapSelectedKeysPushDownUtilsTest.scala b/paimon-spark/paimon-spark-common/src/test/scala/org/apache/paimon/spark/read/MapSelectedKeysPushDownUtilsTest.scala new file mode 100644 index 000000000000..d3019baa5955 --- /dev/null +++ b/paimon-spark/paimon-spark-common/src/test/scala/org/apache/paimon/spark/read/MapSelectedKeysPushDownUtilsTest.scala @@ -0,0 +1,80 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you 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 org.apache.paimon.spark.read + +import org.apache.paimon.data.shredding.MapSelectedKeysMetadataUtils +import org.apache.paimon.types.{DataTypes, MapType, RowType} + +import org.scalatest.funsuite.AnyFunSuite + +import scala.collection.JavaConverters._ + +/** Tests for {@link MapSelectedKeysPushDownUtils}. */ +class MapSelectedKeysPushDownUtilsTest extends AnyFunSuite { + + test("rewrite map field with selected keys") { + val rowType = DataTypes.ROW( + DataTypes.FIELD(0, "id", DataTypes.INT()), + DataTypes.FIELD(1, "attrs", DataTypes.MAP(DataTypes.STRING(), DataTypes.BIGINT().notNull())), + DataTypes.FIELD(2, "tags", DataTypes.MAP(DataTypes.STRING(), DataTypes.STRING())) + ) + + val rewritten = MapSelectedKeysPushDownUtils.rewriteRowType( + rowType, + Map("attrs" -> Seq("key1", "key2")) + ) + + assert(rewritten.getField("id").`type`() == DataTypes.INT()) + assert(rewritten.getField("tags").`type`().isInstanceOf[MapType]) + + val attrs = rewritten.getField("attrs") + assert(attrs.description() == "__PAIMON_MAP_SELECTED_KEYS:key1;key2") + assert(MapSelectedKeysMetadataUtils.isMapSelectedKeysField(attrs)) + assert( + MapSelectedKeysMetadataUtils.selectedKeys(attrs.description()).asScala == Seq("key1", "key2")) + + val attrsRowType = attrs.`type`().asInstanceOf[RowType] + assert(attrsRowType.getFieldNames.asScala == Seq("0", "1")) + assert(attrsRowType.getField("0").`type`() == DataTypes.BIGINT()) + assert(attrsRowType.getField("1").`type`() == DataTypes.BIGINT()) + assert(attrsRowType.getField("0").`type`().isNullable) + assert(attrsRowType.getField("1").`type`().isNullable) + } + + test("ignore nested map field with selected keys") { + val nestedType = DataTypes.ROW( + DataTypes.FIELD(2, "name", DataTypes.STRING()), + DataTypes.FIELD(3, "attrs", DataTypes.MAP(DataTypes.STRING(), DataTypes.DOUBLE())) + ) + val rowType = DataTypes.ROW( + DataTypes.FIELD(0, "id", DataTypes.INT()), + DataTypes.FIELD(1, "profile", nestedType) + ) + + val rewritten = MapSelectedKeysPushDownUtils.rewriteRowType( + rowType, + Map("profile.attrs" -> Seq("key1", "key2")) + ) + + val profile = rewritten.getField("profile").`type`().asInstanceOf[RowType] + assert(profile.getField("name").`type`() == DataTypes.STRING()) + assert(profile.getField("attrs").description() == null) + assert(profile.getField("attrs").`type`().isInstanceOf[MapType]) + } +} diff --git a/paimon-spark/paimon-spark-ut/src/test/scala/org/apache/paimon/spark/sql/PaimonOptimizationTestBase.scala b/paimon-spark/paimon-spark-ut/src/test/scala/org/apache/paimon/spark/sql/PaimonOptimizationTestBase.scala index 178006ff1aeb..0550e1e4e669 100644 --- a/paimon-spark/paimon-spark-ut/src/test/scala/org/apache/paimon/spark/sql/PaimonOptimizationTestBase.scala +++ b/paimon-spark/paimon-spark-ut/src/test/scala/org/apache/paimon/spark/sql/PaimonOptimizationTestBase.scala @@ -19,9 +19,10 @@ package org.apache.paimon.spark.sql import org.apache.paimon.Snapshot.CommitKind -import org.apache.paimon.spark.PaimonSparkTestBase +import org.apache.paimon.spark.{PaimonScan, PaimonSparkTestBase} import org.apache.paimon.spark.catalyst.analysis.expressions.ExpressionHelper import org.apache.paimon.spark.catalyst.optimizer.MergePaimonScalarSubqueries +import org.apache.paimon.spark.catalyst.optimizer.PushDownMapSelectedKeys import org.apache.paimon.spark.execution.TruncatePaimonTableWithFilterExec import org.apache.spark.sql.{DataFrame, PaimonUtils, Row} @@ -29,6 +30,7 @@ import org.apache.spark.sql.catalyst.expressions.{Attribute, CreateNamedStruct, import org.apache.spark.sql.catalyst.plans.logical.{CTERelationDef, LogicalPlan, OneRowRelation, WithCTE} import org.apache.spark.sql.catalyst.rules.RuleExecutor import org.apache.spark.sql.execution.CommandResultExec +import org.apache.spark.sql.execution.datasources.v2.DataSourceV2ScanRelation import org.apache.spark.sql.functions._ import org.junit.jupiter.api.Assertions @@ -118,6 +120,44 @@ abstract class PaimonOptimizationTestBase extends PaimonSparkTestBase with Expre } } + test("Paimon Optimization: map selected-key pushdown only supports shared-shredding map") { + withTable("T") { + spark.sql("CREATE TABLE T (id INT, attrs MAP)") + + val normalMapScan = mapSelectedKeysPaimonScan("SELECT attrs['key1'] FROM T") + Assertions.assertTrue(normalMapScan.pushedMapSelectedKeys.isEmpty) + } + + withTable("T") { + spark.sql("CREATE TABLE T (id INT, profile STRUCT>)") + + val nestedMapScan = mapSelectedKeysPaimonScan("SELECT profile.attrs['key1'] FROM T") + Assertions.assertTrue(nestedMapScan.pushedMapSelectedKeys.isEmpty) + } + + withTable("T") { + spark.sql(""" + |CREATE TABLE T (id INT, attrs MAP) + |TBLPROPERTIES ( + | 'fields.attrs.map.storage-layout' = 'shared-shredding' + |) + |""".stripMargin) + + val sharedShreddingMapScan = + mapSelectedKeysPaimonScan("SELECT attrs['key1'], attrs['key2'] FROM T") + Assertions.assertEquals( + Map("attrs" -> Seq("key1", "key2")), + sharedShreddingMapScan.pushedMapSelectedKeys) + + val delimiterKeyScan = mapSelectedKeysPaimonScan("SELECT attrs['a;b'] FROM T") + Assertions.assertTrue(delimiterKeyScan.pushedMapSelectedKeys.isEmpty) + + val metadataPrefixKeyScan = + mapSelectedKeysPaimonScan("SELECT attrs['__PAIMON_MAP_SELECTED_KEYS:key1'] FROM T") + Assertions.assertTrue(metadataPrefixKeyScan.pushedMapSelectedKeys.isEmpty) + } + } + test(s"Paimon Optimization: optimize metadata only delete") { for (useV2Write <- Seq("true", "false")) { withSparkSQLConf("spark.paimon.write.use-v2-write" -> useV2Write) { @@ -233,6 +273,12 @@ abstract class PaimonOptimizationTestBase extends PaimonSparkTestBase with Expre def extractorExpression(cteIndex: Int, output: Seq[Attribute], fieldIndex: Int): NamedExpression + private def mapSelectedKeysPaimonScan(sqlText: String): PaimonScan = { + PushDownMapSelectedKeys(sql(sqlText).queryExecution.optimizedPlan).collectFirst { + case relation: DataSourceV2ScanRelation => relation.scan.asInstanceOf[PaimonScan] + }.get + } + def checkTruncatePaimonTable(df: DataFrame): Unit = { val plan = df.queryExecution.executedPlan.asInstanceOf[CommandResultExec].commandPhysicalPlan assert(plan.isInstanceOf[TruncatePaimonTableWithFilterExec]) diff --git a/paimon-spark/paimon-spark4-common/src/main/scala/org/apache/paimon/spark/catalyst/optimizer/PushDownMapSelectedKeys.scala b/paimon-spark/paimon-spark4-common/src/main/scala/org/apache/paimon/spark/catalyst/optimizer/PushDownMapSelectedKeys.scala new file mode 100644 index 000000000000..88e6347cee2e --- /dev/null +++ b/paimon-spark/paimon-spark4-common/src/main/scala/org/apache/paimon/spark/catalyst/optimizer/PushDownMapSelectedKeys.scala @@ -0,0 +1,39 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you 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 org.apache.paimon.spark.catalyst.optimizer + +import org.apache.paimon.spark.PaimonScan + +import org.apache.spark.sql.catalyst.expressions.AttributeReference +import org.apache.spark.sql.execution.datasources.v2.DataSourceV2ScanRelation + +object PushDownMapSelectedKeys extends PushDownMapSelectedKeysBase { + + override protected def copyDataSourceV2ScanRelation( + relation: DataSourceV2ScanRelation, + scan: PaimonScan, + output: Seq[AttributeReference]): DataSourceV2ScanRelation = { + DataSourceV2ScanRelation( + relation.relation, + scan, + output, + relation.keyGroupedPartitioning, + relation.ordering) + } +}