exec/java-exec/src/test/java/org/apache/drill/exec/planner/fragment/TestHardAffinityFragmentParallelizer.java - drill - Git at Google

 /*
  * 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.drill.exec.planner.fragment;

 import com.google.common.collect.HashMultiset;
 import com.google.common.collect.ImmutableList;
 import org.apache.drill.categories.PlannerTest;
 import org.apache.drill.exec.physical.EndpointAffinity;
 import org.apache.drill.exec.physical.base.PhysicalOperator;
 import org.apache.drill.exec.proto.CoordinationProtos.DrillbitEndpoint;
 import org.apache.drill.test.BaseTest;
 import org.junit.Test;
 import org.junit.experimental.categories.Category;

 import java.util.Collections;
 import java.util.List;

 import static java.lang.Integer.MAX_VALUE;
 import static org.apache.drill.exec.ExecConstants.SLICE_TARGET_DEFAULT;
 import static org.apache.drill.exec.planner.fragment.HardAffinityFragmentParallelizer.INSTANCE;
 import static org.junit.Assert.assertEquals;
 import static org.junit.Assert.assertTrue;
 import static org.junit.Assert.fail;
 import static org.mockito.Mockito.mock;
 import static org.mockito.Mockito.when;

 @Category(PlannerTest.class)
 public class TestHardAffinityFragmentParallelizer extends BaseTest {

   // Create a set of test endpoints
   private static final DrillbitEndpoint N1_EP1 = newDrillbitEndpoint("node1", 30010);
   private static final DrillbitEndpoint N1_EP2 = newDrillbitEndpoint("node1", 30011);
   private static final DrillbitEndpoint N2_EP1 = newDrillbitEndpoint("node2", 30010);
   private static final DrillbitEndpoint N2_EP2 = newDrillbitEndpoint("node2", 30011);
   private static final DrillbitEndpoint N3_EP1 = newDrillbitEndpoint("node3", 30010);
   private static final DrillbitEndpoint N3_EP2 = newDrillbitEndpoint("node3", 30011);
   private static final DrillbitEndpoint N4_EP2 = newDrillbitEndpoint("node4", 30011);

   private static final DrillbitEndpoint newDrillbitEndpoint(String address, int port) {
     return DrillbitEndpoint.newBuilder().setAddress(address).setControlPort(port).build();
   }

   private static final ParallelizationParameters newParameters(final long threshold, final int maxWidthPerNode,
       final int maxGlobalWidth) {
     return new ParallelizationParameters() {
       @Override
       public long getSliceTarget() {
         return threshold;
       }

       @Override
       public int getMaxWidthPerNode() {
         return maxWidthPerNode;
       }

       @Override
       public int getMaxGlobalWidth() {
         return maxGlobalWidth;
       }

       /**
        * {@link HardAffinityFragmentParallelizer} doesn't use affinity factor.
        * @return
        */
       @Override
       public double getAffinityFactor() {
         return 0.0f;
       }
     };
   }

   private final Wrapper newWrapper(double cost, int minWidth, int maxWidth, List<EndpointAffinity> epAffs) {
     final Fragment fragment = mock(Fragment.class);
     final PhysicalOperator root = mock(PhysicalOperator.class);

     when(fragment.getRoot()).thenReturn(root);

     final Wrapper fragmentWrapper = new Wrapper(fragment, 1);
     final Stats stats = fragmentWrapper.getStats();
     stats.setDistributionAffinity(DistributionAffinity.HARD);
     stats.addCost(cost);
     stats.addMinWidth(minWidth);
     stats.addMaxWidth(maxWidth);
     stats.addEndpointAffinities(epAffs);

     return fragmentWrapper;
   }

   @Test
   public void simpleCase1() throws Exception {
     final Wrapper wrapper = newWrapper(200, 1, 20, Collections.singletonList(new EndpointAffinity(N1_EP1, 1.0, true, MAX_VALUE)));
     INSTANCE.parallelizeFragment(wrapper, newParameters(SLICE_TARGET_DEFAULT, 5, 20), null);

     // Expect the fragment parallelization to be just one because:
     // The cost (200) is below the threshold (SLICE_TARGET_DEFAULT) (which gives width of 200/10000 = ~1) and
     assertEquals(1, wrapper.getWidth());

     final List<DrillbitEndpoint> assignedEps = wrapper.getAssignedEndpoints();
     assertEquals(1, assignedEps.size());
     assertEquals(N1_EP1, assignedEps.get(0));
   }

   @Test
   public void simpleCase2() throws Exception {
     // Set the slice target to 1
     final Wrapper wrapper = newWrapper(200, 1, 20, Collections.singletonList(new EndpointAffinity(N1_EP1, 1.0, true, MAX_VALUE)));
     INSTANCE.parallelizeFragment(wrapper, newParameters(1, 5, 20), null);

     // Expect the fragment parallelization to be 5:
     // 1. the cost (200) is above the threshold (SLICE_TARGET_DEFAULT) (which gives 200/1=200 width) and
     // 2. Max width per node is 5 (limits the width 200 to 5)
     assertEquals(5, wrapper.getWidth());

     final List<DrillbitEndpoint> assignedEps = wrapper.getAssignedEndpoints();
     assertEquals(5, assignedEps.size());
     for (DrillbitEndpoint ep : assignedEps) {
       assertEquals(N1_EP1, ep);
     }
   }

   @Test
   public void multiNodeCluster1() throws Exception {
     final Wrapper wrapper = newWrapper(200, 1, 20,
         ImmutableList.of(
             new EndpointAffinity(N1_EP1, 0.15, true, MAX_VALUE),
             new EndpointAffinity(N1_EP2, 0.15, true, MAX_VALUE),
             new EndpointAffinity(N2_EP1, 0.10, true, MAX_VALUE),
             new EndpointAffinity(N3_EP2, 0.20, true, MAX_VALUE),
             new EndpointAffinity(N4_EP2, 0.20, true, MAX_VALUE)
         ));
     INSTANCE.parallelizeFragment(wrapper, newParameters(SLICE_TARGET_DEFAULT, 5, 20), null);

     // Expect the fragment parallelization to be 5 because:
     // 1. The cost (200) is below the threshold (SLICE_TARGET_DEFAULT) (which gives width of 200/10000 = ~1) and
     // 2. Number of mandoatory node assignments are 5 which overrides the cost based width of 1.
     assertEquals(5, wrapper.getWidth());

     // As there are 5 required eps and the width is 5, everyone gets assigned 1.
     final List<DrillbitEndpoint> assignedEps = wrapper.getAssignedEndpoints();
     assertEquals(5, assignedEps.size());
     assertTrue(assignedEps.contains(N1_EP1));
     assertTrue(assignedEps.contains(N1_EP2));
     assertTrue(assignedEps.contains(N2_EP1));
     assertTrue(assignedEps.contains(N3_EP2));
     assertTrue(assignedEps.contains(N4_EP2));
   }

   @Test
   public void multiNodeCluster2() throws Exception {
     final Wrapper wrapper = newWrapper(200, 1, 20,
         ImmutableList.of(
             new EndpointAffinity(N1_EP2, 0.15, true, MAX_VALUE),
             new EndpointAffinity(N2_EP2, 0.15, true, MAX_VALUE),
             new EndpointAffinity(N3_EP1, 0.10, true, MAX_VALUE),
             new EndpointAffinity(N4_EP2, 0.20, true, MAX_VALUE),
             new EndpointAffinity(N1_EP1, 0.20, true, MAX_VALUE)
         ));
     INSTANCE.parallelizeFragment(wrapper, newParameters(1, 5, 20), null);

     // Expect the fragment parallelization to be 20 because:
     // 1. the cost (200) is above the threshold (SLICE_TARGET_DEFAULT) (which gives 200/1=200 width) and
     // 2. Number of mandatory node assignments are 5 (current width 200 satisfies the requirement)
     // 3. max fragment width is 20 which limits the width
     assertEquals(20, wrapper.getWidth());

     final List<DrillbitEndpoint> assignedEps = wrapper.getAssignedEndpoints();
     assertEquals(20, assignedEps.size());
     final HashMultiset<DrillbitEndpoint> counts = HashMultiset.create();
     for(final DrillbitEndpoint ep : assignedEps) {
       counts.add(ep);
     }
     // Each node gets at max 5.
     assertTrue(counts.count(N1_EP2) <= 5);
     assertTrue(counts.count(N2_EP2) <= 5);
     assertTrue(counts.count(N3_EP1) <= 5);
     assertTrue(counts.count(N4_EP2) <= 5);
     assertTrue(counts.count(N1_EP1) <= 5);
   }

   @Test
   public void multiNodeClusterNegative1() throws Exception {
     final Wrapper wrapper = newWrapper(200, 1, 20,
         ImmutableList.of(
             new EndpointAffinity(N1_EP2, 0.15, true, MAX_VALUE),
             new EndpointAffinity(N2_EP2, 0.15, true, MAX_VALUE),
             new EndpointAffinity(N3_EP1, 0.10, true, MAX_VALUE),
             new EndpointAffinity(N4_EP2, 0.20, true, MAX_VALUE),
             new EndpointAffinity(N1_EP1, 0.20, true, MAX_VALUE)
         ));

     try {
       INSTANCE.parallelizeFragment(wrapper, newParameters(1, 2, 2), null);
       fail("Expected an exception, because max global query width (2) is less than the number of mandatory nodes (5)");
     } catch (Exception e) {
       // ok
     }
   }

   @Test
   public void multiNodeClusterNegative2() throws Exception {
     final Wrapper wrapper = newWrapper(200, 1, 3,
         ImmutableList.of(
             new EndpointAffinity(N1_EP2, 0.15, true, MAX_VALUE),
             new EndpointAffinity(N2_EP2, 0.15, true, MAX_VALUE),
             new EndpointAffinity(N3_EP1, 0.10, true, MAX_VALUE),
             new EndpointAffinity(N4_EP2, 0.20, true, MAX_VALUE),
             new EndpointAffinity(N1_EP1, 0.20, true, MAX_VALUE)
         ));

     try {
       INSTANCE.parallelizeFragment(wrapper, newParameters(1, 2, 2), null);
       fail("Expected an exception, because max fragment width (3) is less than the number of mandatory nodes (5)");
     } catch (Exception e) {
       // ok
     }
   }
 }
	/*
	* 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.drill.exec.planner.fragment;

	import com.google.common.collect.HashMultiset;
	import com.google.common.collect.ImmutableList;
	import org.apache.drill.categories.PlannerTest;
	import org.apache.drill.exec.physical.EndpointAffinity;
	import org.apache.drill.exec.physical.base.PhysicalOperator;
	import org.apache.drill.exec.proto.CoordinationProtos.DrillbitEndpoint;
	import org.apache.drill.test.BaseTest;
	import org.junit.Test;
	import org.junit.experimental.categories.Category;

	import java.util.Collections;
	import java.util.List;

	import static java.lang.Integer.MAX_VALUE;
	import static org.apache.drill.exec.ExecConstants.SLICE_TARGET_DEFAULT;
	import static org.apache.drill.exec.planner.fragment.HardAffinityFragmentParallelizer.INSTANCE;
	import static org.junit.Assert.assertEquals;
	import static org.junit.Assert.assertTrue;
	import static org.junit.Assert.fail;
	import static org.mockito.Mockito.mock;
	import static org.mockito.Mockito.when;

	@Category(PlannerTest.class)
	public class TestHardAffinityFragmentParallelizer extends BaseTest {

	// Create a set of test endpoints
	private static final DrillbitEndpoint N1_EP1 = newDrillbitEndpoint("node1", 30010);
	private static final DrillbitEndpoint N1_EP2 = newDrillbitEndpoint("node1", 30011);
	private static final DrillbitEndpoint N2_EP1 = newDrillbitEndpoint("node2", 30010);
	private static final DrillbitEndpoint N2_EP2 = newDrillbitEndpoint("node2", 30011);
	private static final DrillbitEndpoint N3_EP1 = newDrillbitEndpoint("node3", 30010);
	private static final DrillbitEndpoint N3_EP2 = newDrillbitEndpoint("node3", 30011);
	private static final DrillbitEndpoint N4_EP2 = newDrillbitEndpoint("node4", 30011);

	private static final DrillbitEndpoint newDrillbitEndpoint(String address, int port) {
	return DrillbitEndpoint.newBuilder().setAddress(address).setControlPort(port).build();
	}

	private static final ParallelizationParameters newParameters(final long threshold, final int maxWidthPerNode,
	final int maxGlobalWidth) {
	return new ParallelizationParameters() {
	@Override
	public long getSliceTarget() {
	return threshold;
	}

	@Override
	public int getMaxWidthPerNode() {
	return maxWidthPerNode;
	}

	@Override
	public int getMaxGlobalWidth() {
	return maxGlobalWidth;
	}

	/**
	* {@link HardAffinityFragmentParallelizer} doesn't use affinity factor.
	* @return
	*/
	@Override
	public double getAffinityFactor() {
	return 0.0f;
	}
	};
	}

	private final Wrapper newWrapper(double cost, int minWidth, int maxWidth, List<EndpointAffinity> epAffs) {
	final Fragment fragment = mock(Fragment.class);
	final PhysicalOperator root = mock(PhysicalOperator.class);

	when(fragment.getRoot()).thenReturn(root);

	final Wrapper fragmentWrapper = new Wrapper(fragment, 1);
	final Stats stats = fragmentWrapper.getStats();
	stats.setDistributionAffinity(DistributionAffinity.HARD);
	stats.addCost(cost);
	stats.addMinWidth(minWidth);
	stats.addMaxWidth(maxWidth);
	stats.addEndpointAffinities(epAffs);

	return fragmentWrapper;
	}

	@Test
	public void simpleCase1() throws Exception {
	final Wrapper wrapper = newWrapper(200, 1, 20, Collections.singletonList(new EndpointAffinity(N1_EP1, 1.0, true, MAX_VALUE)));
	INSTANCE.parallelizeFragment(wrapper, newParameters(SLICE_TARGET_DEFAULT, 5, 20), null);

	// Expect the fragment parallelization to be just one because:
	// The cost (200) is below the threshold (SLICE_TARGET_DEFAULT) (which gives width of 200/10000 = ~1) and
	assertEquals(1, wrapper.getWidth());

	final List<DrillbitEndpoint> assignedEps = wrapper.getAssignedEndpoints();
	assertEquals(1, assignedEps.size());
	assertEquals(N1_EP1, assignedEps.get(0));
	}

	@Test
	public void simpleCase2() throws Exception {
	// Set the slice target to 1
	final Wrapper wrapper = newWrapper(200, 1, 20, Collections.singletonList(new EndpointAffinity(N1_EP1, 1.0, true, MAX_VALUE)));
	INSTANCE.parallelizeFragment(wrapper, newParameters(1, 5, 20), null);

	// Expect the fragment parallelization to be 5:
	// 1. the cost (200) is above the threshold (SLICE_TARGET_DEFAULT) (which gives 200/1=200 width) and
	// 2. Max width per node is 5 (limits the width 200 to 5)
	assertEquals(5, wrapper.getWidth());

	final List<DrillbitEndpoint> assignedEps = wrapper.getAssignedEndpoints();
	assertEquals(5, assignedEps.size());
	for (DrillbitEndpoint ep : assignedEps) {
	assertEquals(N1_EP1, ep);
	}
	}

	@Test
	public void multiNodeCluster1() throws Exception {
	final Wrapper wrapper = newWrapper(200, 1, 20,
	ImmutableList.of(
	new EndpointAffinity(N1_EP1, 0.15, true, MAX_VALUE),
	new EndpointAffinity(N1_EP2, 0.15, true, MAX_VALUE),
	new EndpointAffinity(N2_EP1, 0.10, true, MAX_VALUE),
	new EndpointAffinity(N3_EP2, 0.20, true, MAX_VALUE),
	new EndpointAffinity(N4_EP2, 0.20, true, MAX_VALUE)
	));
	INSTANCE.parallelizeFragment(wrapper, newParameters(SLICE_TARGET_DEFAULT, 5, 20), null);

	// Expect the fragment parallelization to be 5 because:
	// 1. The cost (200) is below the threshold (SLICE_TARGET_DEFAULT) (which gives width of 200/10000 = ~1) and
	// 2. Number of mandoatory node assignments are 5 which overrides the cost based width of 1.
	assertEquals(5, wrapper.getWidth());

	// As there are 5 required eps and the width is 5, everyone gets assigned 1.
	final List<DrillbitEndpoint> assignedEps = wrapper.getAssignedEndpoints();
	assertEquals(5, assignedEps.size());
	assertTrue(assignedEps.contains(N1_EP1));
	assertTrue(assignedEps.contains(N1_EP2));
	assertTrue(assignedEps.contains(N2_EP1));
	assertTrue(assignedEps.contains(N3_EP2));
	assertTrue(assignedEps.contains(N4_EP2));
	}

	@Test
	public void multiNodeCluster2() throws Exception {
	final Wrapper wrapper = newWrapper(200, 1, 20,
	ImmutableList.of(
	new EndpointAffinity(N1_EP2, 0.15, true, MAX_VALUE),
	new EndpointAffinity(N2_EP2, 0.15, true, MAX_VALUE),
	new EndpointAffinity(N3_EP1, 0.10, true, MAX_VALUE),
	new EndpointAffinity(N4_EP2, 0.20, true, MAX_VALUE),
	new EndpointAffinity(N1_EP1, 0.20, true, MAX_VALUE)
	));
	INSTANCE.parallelizeFragment(wrapper, newParameters(1, 5, 20), null);

	// Expect the fragment parallelization to be 20 because:
	// 1. the cost (200) is above the threshold (SLICE_TARGET_DEFAULT) (which gives 200/1=200 width) and
	// 2. Number of mandatory node assignments are 5 (current width 200 satisfies the requirement)
	// 3. max fragment width is 20 which limits the width
	assertEquals(20, wrapper.getWidth());

	final List<DrillbitEndpoint> assignedEps = wrapper.getAssignedEndpoints();
	assertEquals(20, assignedEps.size());
	final HashMultiset<DrillbitEndpoint> counts = HashMultiset.create();
	for(final DrillbitEndpoint ep : assignedEps) {
	counts.add(ep);
	}
	// Each node gets at max 5.
	assertTrue(counts.count(N1_EP2) <= 5);
	assertTrue(counts.count(N2_EP2) <= 5);
	assertTrue(counts.count(N3_EP1) <= 5);
	assertTrue(counts.count(N4_EP2) <= 5);
	assertTrue(counts.count(N1_EP1) <= 5);
	}

	@Test
	public void multiNodeClusterNegative1() throws Exception {
	final Wrapper wrapper = newWrapper(200, 1, 20,
	ImmutableList.of(
	new EndpointAffinity(N1_EP2, 0.15, true, MAX_VALUE),
	new EndpointAffinity(N2_EP2, 0.15, true, MAX_VALUE),
	new EndpointAffinity(N3_EP1, 0.10, true, MAX_VALUE),
	new EndpointAffinity(N4_EP2, 0.20, true, MAX_VALUE),
	new EndpointAffinity(N1_EP1, 0.20, true, MAX_VALUE)
	));

	try {
	INSTANCE.parallelizeFragment(wrapper, newParameters(1, 2, 2), null);
	fail("Expected an exception, because max global query width (2) is less than the number of mandatory nodes (5)");
	} catch (Exception e) {
	// ok
	}
	}

	@Test
	public void multiNodeClusterNegative2() throws Exception {
	final Wrapper wrapper = newWrapper(200, 1, 3,
	ImmutableList.of(
	new EndpointAffinity(N1_EP2, 0.15, true, MAX_VALUE),
	new EndpointAffinity(N2_EP2, 0.15, true, MAX_VALUE),
	new EndpointAffinity(N3_EP1, 0.10, true, MAX_VALUE),
	new EndpointAffinity(N4_EP2, 0.20, true, MAX_VALUE),
	new EndpointAffinity(N1_EP1, 0.20, true, MAX_VALUE)
	));

	try {
	INSTANCE.parallelizeFragment(wrapper, newParameters(1, 2, 2), null);
	fail("Expected an exception, because max fragment width (3) is less than the number of mandatory nodes (5)");
	} catch (Exception e) {
	// ok
	}
	}
	}