Support report partitioning to eliminate shuffle exchange

paimon-core/src/main/java/org/apache/paimon/schema/TableSchema.java

@@ -18,6 +18,7 @@
 
 package org.apache.paimon.schema;
 
+import org.apache.paimon.CoreOptions;
 import org.apache.paimon.fs.FileIO;
 import org.apache.paimon.fs.Path;
 import org.apache.paimon.types.DataField;
@@ -66,6 +67,10 @@ public class TableSchema implements Serializable {
 
     private final List<String> primaryKeys;
 
+    private final List<String> bucketKeys;
+
+    private final int numBucket;
+
     private final Map<String, String> options;
 
     private final @Nullable String comment;
@@ -115,8 +120,13 @@ public TableSchema(
         // try to trim to validate primary keys
         trimmedPrimaryKeys();
 
-        // try to validate bucket keys
-        originalBucketKeys();
+        // try to validate and initalize the bucket keys
+        List<String> tmpBucketKeys = originalBucketKeys();
+        if (tmpBucketKeys.isEmpty()) {
+            tmpBucketKeys = trimmedPrimaryKeys();
+        }
+        bucketKeys = tmpBucketKeys;
+        numBucket = CoreOptions.fromMap(options).bucket();
     }
 
     public int version() {
@@ -171,11 +181,11 @@ public Map<String, String> options() {
         return options;
     }
 
+    public int numBuckets() {
+        return numBucket;
+    }
+
     public List<String> bucketKeys() {
-        List<String> bucketKeys = originalBucketKeys();
-        if (bucketKeys.isEmpty()) {
-            bucketKeys = trimmedPrimaryKeys();
-        }
         return bucketKeys;
     }
 

paimon-core/src/main/java/org/apache/paimon/table/AbstractFileStoreTable.java

@@ -150,11 +150,6 @@ public Optional<Statistics> statistics() {
         return Optional.empty();
     }
 
-    @Override
-    public BucketMode bucketMode() {
-        return store().bucketMode();
-    }
-
     @Override
     public Optional<WriteSelector> newWriteSelector() {
         switch (bucketMode()) {

paimon-core/src/main/java/org/apache/paimon/table/BucketSpec.java

@@ -0,0 +1,65 @@
+/*
+ * 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.table;
+
+import java.util.List;
+
+/**
+ * Bucket spec holds all bucket information, we can do plan optimization during table scan.
+ *
+ * <p>If the `bucketMode` is {@link BucketMode#HASH_DYNAMIC}, then `numBucket` is -1;
+ *
+ * @since 0.9
+ */
+public class BucketSpec {
+
+    private final BucketMode bucketMode;
+    private final List<String> bucketKeys;
+    private final int numBuckets;
+
+    public BucketSpec(BucketMode bucketMode, List<String> bucketKeys, int numBuckets) {
+        this.bucketMode = bucketMode;
+        this.bucketKeys = bucketKeys;
+        this.numBuckets = numBuckets;
+    }
+
+    public BucketMode getBucketMode() {
+        return bucketMode;
+    }
+
+    public List<String> getBucketKeys() {
+        return bucketKeys;
+    }
+
+    public int getNumBuckets() {
+        return numBuckets;
+    }
+
+    @Override
+    public String toString() {
+        return "BucketSpec{"
+                + "bucketMode="
+                + bucketMode
+                + ", bucketKeys="
+                + bucketKeys
+                + ", numBuckets="
+                + numBuckets
+                + '}';
+    }
+}

paimon-core/src/main/java/org/apache/paimon/table/DelegatedFileStoreTable.java

@@ -113,11 +113,6 @@ public FileStore<?> store() {
         return wrapped.store();
     }
 
-    @Override
-    public BucketMode bucketMode() {
-        return wrapped.bucketMode();
-    }
-
     @Override
     public CatalogEnvironment catalogEnvironment() {
         return wrapped.catalogEnvironment();

paimon-core/src/main/java/org/apache/paimon/table/FileStoreTable.java

@@ -59,6 +59,14 @@ default List<String> primaryKeys() {
         return schema().primaryKeys();
     }
 
+    default BucketSpec bucketSpec() {
+        return new BucketSpec(bucketMode(), schema().bucketKeys(), schema().numBuckets());
+    }
+
+    default BucketMode bucketMode() {
+        return store().bucketMode();
+    }
+
     @Override
     default Map<String, String> options() {
         return schema().options();
@@ -73,8 +81,6 @@ default Optional<String> comment() {
 
     FileStore<?> store();
 
-    BucketMode bucketMode();
-
     CatalogEnvironment catalogEnvironment();
 
     @Override

paimon-spark/paimon-spark-3.3/src/main/java/org/apache/paimon/spark/catalog/PaimonFunctions.java

@@ -0,0 +1,88 @@
+/*
+ * 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.catalog;
+
+import org.apache.spark.sql.connector.catalog.functions.BoundFunction;
+import org.apache.spark.sql.connector.catalog.functions.UnboundFunction;
+import org.apache.spark.sql.types.DataType;
+import org.apache.spark.sql.types.StructField;
+import org.apache.spark.sql.types.StructType;
+
+import static org.apache.paimon.utils.Preconditions.checkArgument;
+import static org.apache.spark.sql.types.DataTypes.IntegerType;
+
+/**
+ * It should be only used for resolving, e.g., for {@link
+ * org.apache.spark.sql.connector.read.SupportsReportPartitioning}.
+ */
+public class PaimonFunctions {
+    /**
+     * For now, we only support report bucket partitioning for table scan. So the case `SELECT
+     * bucket(10, col)` would fail since we do not implement {@link
+     * org.apache.spark.sql.connector.catalog.functions.ScalarFunction}
+     */
+    public static class BucketFunction implements UnboundFunction {
+        @Override
+        public BoundFunction bind(StructType inputType) {
+            if (inputType.size() != 2) {
+                throw new UnsupportedOperationException(
+                        "Wrong number of inputs (expected numBuckets and value)");
+            }
+
+            StructField numBucket = inputType.fields()[0];
+            StructField bucketField = inputType.fields()[1];
+            checkArgument(
+                    numBucket.dataType() == IntegerType,
+                    "bucket number field must be integer type");
+
+            return new BoundFunction() {
+                @Override
+                public DataType[] inputTypes() {
+                    return new DataType[] {IntegerType, bucketField.dataType()};
+                }
+
+                @Override
+                public DataType resultType() {
+                    return IntegerType;
+                }
+
+                @Override
+                public String name() {
+                    return "bucket";
+                }
+
+                @Override
+                public String canonicalName() {
+                    // We have to override this method to make it support canonical equivalent
+                    return "paimon.bucket(" + bucketField.dataType().catalogString() + ", int)";
+                }
+            };
+        }
+
+        @Override
+        public String description() {
+            return name();
+        }
+
+        @Override
+        public String name() {
+            return "bucket";
+        }
+    }
+}

paimon-spark/paimon-spark-3.3/src/main/java/org/apache/paimon/spark/catalog/SparkBaseCatalog.java

@@ -0,0 +1,97 @@
+/*
+ * 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.catalog;
+
+import org.apache.paimon.catalog.Catalog;
+import org.apache.paimon.spark.SparkProcedures;
+import org.apache.paimon.spark.SparkSource;
+import org.apache.paimon.spark.analysis.NoSuchProcedureException;
+import org.apache.paimon.spark.procedure.Procedure;
+import org.apache.paimon.spark.procedure.ProcedureBuilder;
+
+import org.apache.paimon.shade.guava30.com.google.common.collect.ImmutableMap;
+
+import org.apache.spark.sql.catalyst.analysis.NoSuchFunctionException;
+import org.apache.spark.sql.catalyst.analysis.NoSuchNamespaceException;
+import org.apache.spark.sql.connector.catalog.FunctionCatalog;
+import org.apache.spark.sql.connector.catalog.Identifier;
+import org.apache.spark.sql.connector.catalog.SupportsNamespaces;
+import org.apache.spark.sql.connector.catalog.TableCatalog;
+import org.apache.spark.sql.connector.catalog.functions.UnboundFunction;
+
+import java.util.Arrays;
+import java.util.Map;
+
+import scala.Option;
+
+/** Spark base catalog. */
+public abstract class SparkBaseCatalog
+        implements TableCatalog,
+                FunctionCatalog,
+                SupportsNamespaces,
+                ProcedureCatalog,
+                WithPaimonCatalog {
+
+    protected String catalogName;
+
+    @Override
+    public String name() {
+        return catalogName;
+    }
+
+    @Override
+    public Procedure loadProcedure(Identifier identifier) throws NoSuchProcedureException {
+        if (Catalog.SYSTEM_DATABASE_NAME.equals(identifier.namespace()[0])) {
+            ProcedureBuilder builder = SparkProcedures.newBuilder(identifier.name());
+            if (builder != null) {
+                return builder.withTableCatalog(this).build();
+            }
+        }
+        throw new NoSuchProcedureException(identifier);
+    }
+
+    public boolean usePaimon(String provider) {
+        return provider == null || SparkSource.NAME().equalsIgnoreCase(provider);
+    }
+
+    // --------------------- Function Catalog Methods ----------------------------
+    private static final Map<String, UnboundFunction> FUNCTIONS =
+            ImmutableMap.of("bucket", new PaimonFunctions.BucketFunction());
+
+    @Override
+    public UnboundFunction loadFunction(Identifier ident) throws NoSuchFunctionException {
+        UnboundFunction func = FUNCTIONS.get(ident.name());
+        if (func == null) {
+            throw new NoSuchFunctionException(
+                    "Function " + ident + " is not a paimon function", Option.empty());
+        }
+        return func;
+    }
+
+    @Override
+    public Identifier[] listFunctions(String[] namespace) throws NoSuchNamespaceException {
+        if (namespace.length != 0) {
+            throw new NoSuchNamespaceException(
+                    "Namespace " + Arrays.toString(namespace) + " is not valid", Option.empty());
+        }
+        return FUNCTIONS.keySet().stream()
+                .map(name -> Identifier.of(namespace, name))
+                .toArray(Identifier[]::new);
+    }
+}

paimon-spark/paimon-spark-3.3/src/main/scala/org/apache/paimon/spark/PaimonBucketedInputPartition.scala

@@ -0,0 +1,32 @@
+/*
+ * 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
+
+import org.apache.paimon.table.source.Split
+
+import org.apache.spark.sql.catalyst.InternalRow
+import org.apache.spark.sql.catalyst.expressions.GenericInternalRow
+import org.apache.spark.sql.connector.read.{HasPartitionKey, SupportsReportPartitioning}
+
+/** Bucketed input partition should work with [[SupportsReportPartitioning]] together. */
+case class PaimonBucketedInputPartition(splits: Seq[Split], bucket: Int)
+  extends PaimonInputPartition
+  with HasPartitionKey {
+  override def partitionKey(): InternalRow = new GenericInternalRow(Array(bucket.asInstanceOf[Any]))
+}