/*
* 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 com.facebook.presto.sql.planner.optimizations;
import com.facebook.presto.spi.ConstantProperty;
import com.facebook.presto.spi.GroupingProperty;
import com.facebook.presto.spi.SortingProperty;
import com.facebook.presto.spi.block.SortOrder;
import com.facebook.presto.sql.planner.Symbol;
import com.facebook.presto.sql.planner.optimizations.ActualProperties.Global;
import com.google.common.collect.ImmutableList;
import com.google.common.collect.ImmutableSet;
import com.google.common.collect.Lists;
import org.testng.annotations.Test;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.Collections;
import java.util.Comparator;
import java.util.List;
import java.util.Optional;
import static com.facebook.presto.spi.block.SortOrder.ASC_NULLS_FIRST;
import static com.facebook.presto.sql.planner.SystemPartitioningHandle.FIXED_HASH_DISTRIBUTION;
import static com.facebook.presto.sql.planner.optimizations.ActualProperties.Global.arbitraryPartition;
import static com.facebook.presto.sql.planner.optimizations.ActualProperties.Global.partitionedOn;
import static com.facebook.presto.sql.planner.optimizations.ActualProperties.Global.singleStreamPartition;
import static com.facebook.presto.sql.planner.optimizations.ActualProperties.builder;
import static com.facebook.presto.sql.planner.optimizations.AddExchanges.streamingExecutionPreference;
import static com.google.common.collect.ImmutableList.toImmutableList;
import static org.testng.Assert.assertEquals;
public class TestAddExchanges
{
@Test
public void testPickLayoutAnyPreference()
throws Exception
{
Comparator<ActualProperties> preference = streamingExecutionPreference(PreferredProperties.any());
List<ActualProperties> input = ImmutableList.<ActualProperties>builder()
.add(builder()
.global(streamPartitionedOn("a", "b"))
.build())
.add(builder()
.global(singleStreamPartition())
.build())
.add(builder()
.global(arbitraryPartition())
.local(ImmutableList.of(grouped("a", "b")))
.build())
.add(builder()
.global(arbitraryPartition())
.build())
.add(builder()
.global(hashDistributedOn("a"))
.build())
.add(builder()
.global(singleStream())
.local(ImmutableList.of(constant("a"), sorted("b", ASC_NULLS_FIRST)))
.build())
.add(builder()
.global(singleStreamPartition())
.local(ImmutableList.of(sorted("a", ASC_NULLS_FIRST)))
.build())
.build();
// Given no preferences, the original input order should be maintained
assertEquals(stableSort(input, preference), input);
}
@Test
public void testPickLayoutPartitionedPreference()
throws Exception
{
Comparator<ActualProperties> preference = streamingExecutionPreference(PreferredProperties.distributed());
List<ActualProperties> input = ImmutableList.<ActualProperties>builder()
.add(builder()
.global(streamPartitionedOn("a"))
.build())
.add(builder()
.global(singleStreamPartition())
.build())
.add(builder()
.global(arbitraryPartition())
.local(ImmutableList.of(grouped("a", "b")))
.build())
.add(builder()
.global(arbitraryPartition())
.build())
.add(builder()
.global(hashDistributedOn("a"))
.build())
.add(builder()
.global(singleStream())
.local(ImmutableList.of(constant("a"), sorted("b", ASC_NULLS_FIRST)))
.build())
.add(builder()
.global(singleStreamPartition())
.local(ImmutableList.of(sorted("a", ASC_NULLS_FIRST)))
.build())
.build();
List<ActualProperties> expected = ImmutableList.<ActualProperties>builder()
.add(builder()
.global(streamPartitionedOn("a"))
.build())
.add(builder()
.global(arbitraryPartition())
.local(ImmutableList.of(grouped("a", "b")))
.build())
.add(builder()
.global(arbitraryPartition())
.build())
.add(builder()
.global(hashDistributedOn("a"))
.build())
.add(builder()
.global(singleStream())
.local(ImmutableList.of(constant("a"), sorted("b", ASC_NULLS_FIRST)))
.build())
.add(builder()
.global(singleStreamPartition())
.build())
.add(builder()
.global(singleStreamPartition())
.local(ImmutableList.of(sorted("a", ASC_NULLS_FIRST)))
.build())
.build();
assertEquals(stableSort(input, preference), expected);
}
@Test
public void testPickLayoutUnpartitionedPreference()
throws Exception
{
Comparator<ActualProperties> preference = streamingExecutionPreference(PreferredProperties.undistributed());
List<ActualProperties> input = ImmutableList.<ActualProperties>builder()
.add(builder()
.global(streamPartitionedOn("a"))
.build())
.add(builder()
.global(singleStreamPartition())
.build())
.add(builder()
.global(arbitraryPartition())
.local(ImmutableList.of(grouped("a", "b")))
.build())
.add(builder()
.global(arbitraryPartition())
.build())
.add(builder()
.global(hashDistributedOn("a"))
.build())
.add(builder()
.global(singleStream())
.local(ImmutableList.of(constant("a"), sorted("b", ASC_NULLS_FIRST)))
.build())
.add(builder()
.global(singleStreamPartition())
.local(ImmutableList.of(sorted("a", ASC_NULLS_FIRST)))
.build())
.build();
List<ActualProperties> expected = ImmutableList.<ActualProperties>builder()
.add(builder()
.global(singleStreamPartition())
.build())
.add(builder()
.global(singleStreamPartition())
.local(ImmutableList.of(sorted("a", ASC_NULLS_FIRST)))
.build())
.add(builder()
.global(streamPartitionedOn("a"))
.build())
.add(builder()
.global(arbitraryPartition())
.local(ImmutableList.of(grouped("a", "b")))
.build())
.add(builder()
.global(arbitraryPartition())
.build())
.add(builder()
.global(hashDistributedOn("a"))
.build())
.add(builder()
.global(singleStream())
.local(ImmutableList.of(constant("a"), sorted("b", ASC_NULLS_FIRST)))
.build())
.build();
assertEquals(stableSort(input, preference), expected);
}
@Test
public void testPickLayoutPartitionedOnSingle()
throws Exception
{
Comparator<ActualProperties> preference = streamingExecutionPreference(
PreferredProperties.partitioned(ImmutableSet.of(symbol("a"))));
List<ActualProperties> input = ImmutableList.<ActualProperties>builder()
.add(builder()
.global(streamPartitionedOn("a"))
.build())
.add(builder()
.global(singleStreamPartition())
.build())
.add(builder()
.global(arbitraryPartition())
.local(ImmutableList.of(grouped("a", "b")))
.build())
.add(builder()
.global(arbitraryPartition())
.build())
.add(builder()
.global(hashDistributedOn("a"))
.build())
.add(builder()
.global(singleStream())
.local(ImmutableList.of(constant("a"), sorted("b", ASC_NULLS_FIRST)))
.build())
.add(builder()
.global(singleStreamPartition())
.local(ImmutableList.of(sorted("a", ASC_NULLS_FIRST)))
.build())
.build();
List<ActualProperties> expected = ImmutableList.<ActualProperties>builder()
.add(builder()
.global(streamPartitionedOn("a"))
.build())
.add(builder()
.global(singleStreamPartition())
.build())
.add(builder()
.global(hashDistributedOn("a"))
.build())
.add(builder()
.global(singleStream())
.local(ImmutableList.of(constant("a"), sorted("b", ASC_NULLS_FIRST)))
.build())
.add(builder()
.global(singleStreamPartition())
.local(ImmutableList.of(sorted("a", ASC_NULLS_FIRST)))
.build())
.add(builder()
.global(arbitraryPartition())
.local(ImmutableList.of(grouped("a", "b")))
.build())
.add(builder()
.global(arbitraryPartition())
.build())
.build();
assertEquals(stableSort(input, preference), expected);
}
@Test
public void testPickLayoutPartitionedOnMultiple()
throws Exception
{
Comparator<ActualProperties> preference = streamingExecutionPreference(
PreferredProperties.partitioned(ImmutableSet.of(symbol("a"), symbol("b"))));
List<ActualProperties> input = ImmutableList.<ActualProperties>builder()
.add(builder()
.global(streamPartitionedOn("a"))
.build())
.add(builder()
.global(singleStreamPartition())
.build())
.add(builder()
.global(arbitraryPartition())
.local(ImmutableList.of(grouped("a", "b")))
.build())
.add(builder()
.global(arbitraryPartition())
.build())
.add(builder()
.global(hashDistributedOn("a"))
.build())
.add(builder()
.global(singleStream())
.local(ImmutableList.of(constant("a"), sorted("b", ASC_NULLS_FIRST)))
.build())
.add(builder()
.global(singleStreamPartition())
.local(ImmutableList.of(sorted("a", ASC_NULLS_FIRST)))
.build())
.add(builder()
.global(hashDistributedOn("a"))
.build())
.build();
List<ActualProperties> expected = ImmutableList.<ActualProperties>builder()
.add(builder()
.global(streamPartitionedOn("a"))
.build())
.add(builder()
.global(singleStreamPartition())
.build())
.add(builder()
.global(hashDistributedOn("a"))
.build())
.add(builder()
.global(singleStream())
.local(ImmutableList.of(constant("a"), sorted("b", ASC_NULLS_FIRST)))
.build())
.add(builder()
.global(singleStreamPartition())
.local(ImmutableList.of(sorted("a", ASC_NULLS_FIRST)))
.build())
.add(builder()
.global(hashDistributedOn("a"))
.build())
.add(builder()
.global(arbitraryPartition())
.local(ImmutableList.of(grouped("a", "b")))
.build())
.add(builder()
.global(arbitraryPartition())
.build())
.build();
assertEquals(stableSort(input, preference), expected);
}
@Test
public void testPickLayoutGrouped()
throws Exception
{
Comparator<ActualProperties> preference = streamingExecutionPreference
(PreferredProperties.local(ImmutableList.of(grouped("a"))));
List<ActualProperties> input = ImmutableList.<ActualProperties>builder()
.add(builder()
.global(streamPartitionedOn("a"))
.build())
.add(builder()
.global(singleStreamPartition())
.build())
.add(builder()
.global(arbitraryPartition())
.local(ImmutableList.of(grouped("a", "b")))
.build())
.add(builder()
.global(arbitraryPartition())
.build())
.add(builder()
.global(hashDistributedOn("a"))
.build())
.add(builder()
.global(singleStream())
.local(ImmutableList.of(constant("a"), sorted("b", ASC_NULLS_FIRST)))
.build())
.add(builder()
.global(singleStreamPartition())
.local(ImmutableList.of(sorted("a", ASC_NULLS_FIRST)))
.build())
.add(builder()
.global(hashDistributedOn("a"))
.build())
.build();
List<ActualProperties> expected = ImmutableList.<ActualProperties>builder()
.add(builder()
.global(singleStream())
.local(ImmutableList.of(constant("a"), sorted("b", ASC_NULLS_FIRST)))
.build())
.add(builder()
.global(singleStreamPartition())
.local(ImmutableList.of(sorted("a", ASC_NULLS_FIRST)))
.build())
.add(builder()
.global(streamPartitionedOn("a"))
.build())
.add(builder()
.global(singleStreamPartition())
.build())
.add(builder()
.global(arbitraryPartition())
.local(ImmutableList.of(grouped("a", "b")))
.build())
.add(builder()
.global(arbitraryPartition())
.build())
.add(builder()
.global(hashDistributedOn("a"))
.build())
.add(builder()
.global(hashDistributedOn("a"))
.build())
.build();
assertEquals(stableSort(input, preference), expected);
}
@Test
public void testPickLayoutGroupedMultiple()
throws Exception
{
Comparator<ActualProperties> preference = streamingExecutionPreference
(PreferredProperties.local(ImmutableList.of(grouped("a", "b"))));
List<ActualProperties> input = ImmutableList.<ActualProperties>builder()
.add(builder()
.global(streamPartitionedOn("a"))
.build())
.add(builder()
.global(singleStreamPartition())
.build())
.add(builder()
.global(arbitraryPartition())
.local(ImmutableList.of(grouped("a", "b")))
.build())
.add(builder()
.global(arbitraryPartition())
.build())
.add(builder()
.global(hashDistributedOn("a"))
.build())
.add(builder()
.global(singleStream())
.local(ImmutableList.of(constant("a"), sorted("b", ASC_NULLS_FIRST)))
.build())
.add(builder()
.global(singleStreamPartition())
.local(ImmutableList.of(sorted("a", ASC_NULLS_FIRST)))
.build())
.add(builder()
.global(hashDistributedOn("a"))
.build())
.build();
List<ActualProperties> expected = ImmutableList.<ActualProperties>builder()
.add(builder()
.global(arbitraryPartition())
.local(ImmutableList.of(grouped("a", "b")))
.build())
.add(builder()
.global(singleStream())
.local(ImmutableList.of(constant("a"), sorted("b", ASC_NULLS_FIRST)))
.build())
.add(builder()
.global(singleStreamPartition())
.local(ImmutableList.of(sorted("a", ASC_NULLS_FIRST)))
.build())
.add(builder()
.global(streamPartitionedOn("a"))
.build())
.add(builder()
.global(singleStreamPartition())
.build())
.add(builder()
.global(arbitraryPartition())
.build())
.add(builder()
.global(hashDistributedOn("a"))
.build())
.add(builder()
.global(hashDistributedOn("a"))
.build())
.build();
assertEquals(stableSort(input, preference), expected);
}
@Test
public void testPickLayoutGroupedMultipleProperties()
throws Exception
{
Comparator<ActualProperties> preference = streamingExecutionPreference
(PreferredProperties.local(ImmutableList.of(grouped("a"), grouped("b"))));
List<ActualProperties> input = ImmutableList.<ActualProperties>builder()
.add(builder()
.global(streamPartitionedOn("a"))
.build())
.add(builder()
.global(singleStreamPartition())
.build())
.add(builder()
.global(arbitraryPartition())
.local(ImmutableList.of(grouped("a", "b")))
.build())
.add(builder()
.global(arbitraryPartition())
.build())
.add(builder()
.global(hashDistributedOn("a"))
.build())
.add(builder()
.global(singleStream())
.local(ImmutableList.of(constant("a"), sorted("b", ASC_NULLS_FIRST)))
.build())
.add(builder()
.global(singleStreamPartition())
.local(ImmutableList.of(sorted("a", ASC_NULLS_FIRST)))
.build())
.add(builder()
.global(hashDistributedOn("a"))
.build())
.build();
List<ActualProperties> expected = ImmutableList.<ActualProperties>builder()
.add(builder()
.global(singleStream())
.local(ImmutableList.of(constant("a"), sorted("b", ASC_NULLS_FIRST)))
.build())
.add(builder()
.global(singleStreamPartition())
.local(ImmutableList.of(sorted("a", ASC_NULLS_FIRST)))
.build())
.add(builder()
.global(streamPartitionedOn("a"))
.build())
.add(builder()
.global(singleStreamPartition())
.build())
.add(builder()
.global(arbitraryPartition())
.local(ImmutableList.of(grouped("a", "b")))
.build())
.add(builder()
.global(arbitraryPartition())
.build())
.add(builder()
.global(hashDistributedOn("a"))
.build())
.add(builder()
.global(hashDistributedOn("a"))
.build())
.build();
assertEquals(stableSort(input, preference), expected);
}
@Test
public void testPickLayoutGroupedWithSort()
throws Exception
{
Comparator<ActualProperties> preference = streamingExecutionPreference
(PreferredProperties.local(ImmutableList.of(grouped("a"), sorted("b", ASC_NULLS_FIRST))));
List<ActualProperties> input = ImmutableList.<ActualProperties>builder()
.add(builder()
.global(streamPartitionedOn("a"))
.build())
.add(builder()
.global(singleStreamPartition())
.build())
.add(builder()
.global(arbitraryPartition())
.local(ImmutableList.of(grouped("a", "b")))
.build())
.add(builder()
.global(arbitraryPartition())
.build())
.add(builder()
.global(hashDistributedOn("a"))
.build())
.add(builder()
.global(singleStream())
.local(ImmutableList.of(constant("a"), sorted("b", ASC_NULLS_FIRST)))
.build())
.add(builder()
.global(singleStreamPartition())
.local(ImmutableList.of(sorted("a", ASC_NULLS_FIRST)))
.build())
.add(builder()
.global(hashDistributedOn("a"))
.build())
.build();
List<ActualProperties> expected = ImmutableList.<ActualProperties>builder()
.add(builder()
.global(singleStream())
.local(ImmutableList.of(constant("a"), sorted("b", ASC_NULLS_FIRST)))
.build())
.add(builder()
.global(singleStreamPartition())
.local(ImmutableList.of(sorted("a", ASC_NULLS_FIRST)))
.build())
.add(builder()
.global(streamPartitionedOn("a"))
.build())
.add(builder()
.global(singleStreamPartition())
.build())
.add(builder()
.global(arbitraryPartition())
.local(ImmutableList.of(grouped("a", "b")))
.build())
.add(builder()
.global(arbitraryPartition())
.build())
.add(builder()
.global(hashDistributedOn("a"))
.build())
.add(builder()
.global(hashDistributedOn("a"))
.build())
.build();
assertEquals(stableSort(input, preference), expected);
}
@Test
public void testPickLayoutUnpartitionedWithGroupAndSort()
throws Exception
{
Comparator<ActualProperties> preference = streamingExecutionPreference
(PreferredProperties.undistributedWithLocal(ImmutableList.of(grouped("a"), sorted("b", ASC_NULLS_FIRST))));
List<ActualProperties> input = ImmutableList.<ActualProperties>builder()
.add(builder()
.global(streamPartitionedOn("a"))
.build())
.add(builder()
.global(singleStreamPartition())
.build())
.add(builder()
.global(arbitraryPartition())
.local(ImmutableList.of(grouped("a", "b")))
.build())
.add(builder()
.global(arbitraryPartition())
.build())
.add(builder()
.global(hashDistributedOn("a"))
.build())
.add(builder()
.global(singleStream())
.local(ImmutableList.of(constant("a"), sorted("b", ASC_NULLS_FIRST)))
.build())
.add(builder()
.global(singleStreamPartition())
.local(ImmutableList.of(sorted("a", ASC_NULLS_FIRST)))
.build())
.add(builder()
.global(hashDistributedOn("a"))
.build())
.build();
List<ActualProperties> expected = ImmutableList.<ActualProperties>builder()
.add(builder()
.global(singleStreamPartition())
.local(ImmutableList.of(sorted("a", ASC_NULLS_FIRST)))
.build())
.add(builder()
.global(singleStream())
.local(ImmutableList.of(constant("a"), sorted("b", ASC_NULLS_FIRST)))
.build())
.add(builder()
.global(singleStreamPartition())
.build())
.add(builder()
.global(streamPartitionedOn("a"))
.build())
.add(builder()
.global(arbitraryPartition())
.local(ImmutableList.of(grouped("a", "b")))
.build())
.add(builder()
.global(arbitraryPartition())
.build())
.add(builder()
.global(hashDistributedOn("a"))
.build())
.add(builder()
.global(hashDistributedOn("a"))
.build())
.build();
assertEquals(stableSort(input, preference), expected);
}
@Test
public void testPickLayoutPartitionedWithGroup()
throws Exception
{
Comparator<ActualProperties> preference = streamingExecutionPreference
(PreferredProperties.partitionedWithLocal(
ImmutableSet.of(symbol("a")),
ImmutableList.of(grouped("a"))));
List<ActualProperties> input = ImmutableList.<ActualProperties>builder()
.add(builder()
.global(streamPartitionedOn("a"))
.build())
.add(builder()
.global(singleStreamPartition())
.build())
.add(builder()
.global(arbitraryPartition())
.local(ImmutableList.of(grouped("a", "b")))
.build())
.add(builder()
.global(arbitraryPartition())
.build())
.add(builder()
.global(hashDistributedOn("a"))
.build())
.add(builder()
.global(singleStream())
.local(ImmutableList.of(constant("a"), sorted("b", ASC_NULLS_FIRST)))
.build())
.add(builder()
.global(singleStreamPartition())
.local(ImmutableList.of(sorted("a", ASC_NULLS_FIRST)))
.build())
.add(builder()
.global(hashDistributedOn("a"))
.build())
.build();
List<ActualProperties> expected = ImmutableList.<ActualProperties>builder()
.add(builder()
.global(singleStream())
.local(ImmutableList.of(constant("a"), sorted("b", ASC_NULLS_FIRST)))
.build())
.add(builder()
.global(singleStreamPartition())
.local(ImmutableList.of(sorted("a", ASC_NULLS_FIRST)))
.build())
.add(builder()
.global(streamPartitionedOn("a"))
.build())
.add(builder()
.global(singleStreamPartition())
.build())
.add(builder()
.global(hashDistributedOn("a"))
.build())
.add(builder()
.global(hashDistributedOn("a"))
.build())
.add(builder()
.global(arbitraryPartition())
.local(ImmutableList.of(grouped("a", "b")))
.build())
.add(builder()
.global(arbitraryPartition())
.build())
.build();
assertEquals(stableSort(input, preference), expected);
}
private static <T> List<T> stableSort(List<T> list, Comparator<T> comparator)
{
ArrayList<T> copy = Lists.newArrayList(list);
Collections.sort(copy, comparator);
return copy;
}
private static Global hashDistributedOn(String... columnNames)
{
return partitionedOn(FIXED_HASH_DISTRIBUTION, arguments(columnNames), Optional.of(arguments(columnNames)));
}
public static Global singleStream()
{
return Global.streamPartitionedOn(ImmutableList.of());
}
private static Global streamPartitionedOn(String... columnNames)
{
return Global.streamPartitionedOn(arguments(columnNames));
}
private static ConstantProperty<Symbol> constant(String column)
{
return new ConstantProperty<>(symbol(column));
}
private static GroupingProperty<Symbol> grouped(String... columns)
{
return new GroupingProperty<>(Lists.transform(Arrays.asList(columns), Symbol::new));
}
private static SortingProperty<Symbol> sorted(String column, SortOrder order)
{
return new SortingProperty<>(symbol(column), order);
}
private static Symbol symbol(String name)
{
return new Symbol(name);
}
private static List<Symbol> arguments(String[] columnNames)
{
return Arrays.asList(columnNames).stream()
.map(Symbol::new)
.collect(toImmutableList());
}
}