/*
 * 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.trino.sql.planner.iterative.rule;

import io.trino.Session;
import io.trino.matching.Captures;
import io.trino.matching.Pattern;
import io.trino.sql.planner.iterative.Rule;
import io.trino.sql.planner.plan.TableWriterNode;

import java.util.Optional;

import static io.trino.SystemSessionProperties.getPreferredWritePartitioningMinNumberOfPartitions;
import static io.trino.SystemSessionProperties.isUsePreferredWritePartitioning;
import static io.trino.cost.AggregationStatsRule.getRowsCount;
import static io.trino.sql.planner.plan.Patterns.tableWriterNode;
import static java.lang.Double.isNaN;

public class PreferWritePartitioning
        implements Rule<TableWriterNode>
{
    @Override
    public Pattern<TableWriterNode> getPattern()
    {
        return tableWriterNode();
    }

    @Override
    public boolean isEnabled(Session session)
    {
        return isUsePreferredWritePartitioning(session);
    }

    @Override
    public Result apply(TableWriterNode node, Captures captures, Context context)
    {
        if (node.getPreferredPartitioningScheme().isEmpty()) {
            return Result.empty();
        }

        int minimumNumberOfPartitions = getPreferredWritePartitioningMinNumberOfPartitions(context.getSession());
        if (minimumNumberOfPartitions <= 1) {
            // Force 'preferred write partitioning' even if stats are missing or broken
            return enable(node);
        }

        double expectedNumberOfPartitions = getRowsCount(
                context.getStatsProvider().getStats(node.getSource()),
                node.getPreferredPartitioningScheme().get().getPartitioning().getColumns());

        if (isNaN(expectedNumberOfPartitions) || expectedNumberOfPartitions < minimumNumberOfPartitions) {
            return Result.empty();
        }

        return enable(node);
    }

    private Result enable(TableWriterNode node)
    {
        return Result.ofPlanNode(new TableWriterNode(
                node.getId(),
                node.getSource(),
                node.getTarget(),
                node.getRowCountSymbol(),
                node.getFragmentSymbol(),
                node.getColumns(),
                node.getColumnNames(),
                node.getNotNullColumnSymbols(),
                node.getPreferredPartitioningScheme(),
                Optional.empty(),
                node.getStatisticsAggregation(),
                node.getStatisticsAggregationDescriptor()));
    }
}

@@ -1245,16 +1245,208 @@
 Node%3E%0A%7B%0A
+    public static final Pattern%3CTableWriterNode%3E WRITER_NODE_WITH_PREFERRED_PARTITIONING = tableWriterNode()%0A            .matching(node -%3E node.getPreferredPartitioningScheme().isPresent());%0A%0A
     @Ove
@@ -1521,33 +1521,55 @@
  return 
-tableWriterNode()
+WRITER_NODE_WITH_PREFERRED_PARTITIONING
 ;%0A    %7D%0A
@@ -1805,117 +1805,8 @@
   %7B%0A
-        if (node.getPreferredPartitioningScheme().isEmpty()) %7B%0A            return Result.empty();%0A        %7D%0A%0A
     

Make optimizer predicate execute earlier