[spark] prepare MAP selected-key pushdown read type

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.
+ *
+ * <p>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<String> keys) {
+        List<String> 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<String> keys) {
+        checkArgument(rowType instanceof RowType, "Selected-key MAP read type must be ROW.");
+        return field.newType(rowType).newDescription(buildMapSelectedKeysMetadata(keys));
+    }
+
+    public static List<String> 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<String> keys = new ArrayList<>(parts.length);
+        for (String part : parts) {
+            keys.add(part);
+        }
+        return keys;
+    }
+
+    private static List<String> normalizeKeys(List<String> keys) {
+        checkNotNull(keys, "Selected keys must not be null.");
+        checkArgument(!keys.isEmpty(), "Selected keys must not be empty.");
+
+        Set<String> 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);
+    }
+}

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.");
+    }
+}

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) {}

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)
+  }
+}

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 {

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)
+  }
+}

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)
+  }
+}

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)
+  }
+}