/**
* 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.hadoop.hive.ql.io;
import static org.mockito.Matchers.anyObject;
import static org.mockito.Mockito.inOrder;
import static org.mockito.Mockito.mock;
import static org.mockito.Mockito.verify;
import static org.mockito.Mockito.when;
import java.io.ByteArrayOutputStream;
import java.io.IOException;
import java.util.LinkedHashMap;
import junit.framework.Assert;
import junit.framework.TestCase;
import org.apache.hadoop.fs.Path;
import org.apache.hadoop.hive.ql.exec.Utilities;
import org.apache.hadoop.hive.ql.io.HiveInputFormat.HiveInputSplit;
import org.apache.hadoop.hive.ql.plan.MapredWork;
import org.apache.hadoop.hive.ql.plan.PartitionDesc;
import org.apache.hadoop.hive.ql.plan.TableDesc;
import org.apache.hadoop.hive.ql.udf.generic.GenericUDFOPEqual;
import org.apache.hadoop.hive.ql.udf.generic.GenericUDFOPEqualOrGreaterThan;
import org.apache.hadoop.hive.ql.udf.generic.GenericUDFOPEqualOrLessThan;
import org.apache.hadoop.hive.ql.udf.generic.GenericUDFOPGreaterThan;
import org.apache.hadoop.hive.ql.udf.generic.GenericUDFOPLessThan;
import org.apache.hadoop.hive.serde2.columnar.BytesRefArrayWritable;
import org.apache.hadoop.io.LongWritable;
import org.apache.hadoop.io.Writable;
import org.apache.hadoop.io.WritableComparable;
import org.apache.hadoop.mapred.JobConf;
import org.apache.hadoop.mapred.RecordReader;
import org.mockito.InOrder;
/**
* TestHiveBinarySearchRecordReader.
*
*/
public class TestHiveBinarySearchRecordReader extends TestCase {
private RCFileRecordReader rcfReader;
private JobConf conf;
private TestHiveInputSplit hiveSplit;
private HiveContextAwareRecordReader hbsReader;
private IOContext ioContext;
private static class TestHiveInputSplit extends HiveInputSplit {
@Override
public long getStart() {
return 0;
}
@Override
public long getLength() {
return 100;
}
@Override
public Path getPath() {
return new Path("/");
}
}
private static class TestHiveRecordReader<K extends WritableComparable, V extends Writable>
extends HiveContextAwareRecordReader<K, V> {
public TestHiveRecordReader(RecordReader recordReader, JobConf conf) throws IOException {
super(recordReader, conf);
}
@Override
public K createKey() {
return null;
}
@Override
public V createValue() {
return null;
}
@Override
public long getPos() throws IOException {
return 0;
}
@Override
public boolean doNext(K key, V value) throws IOException {
return super.doNext(key, value);
}
@Override
public void doClose() throws IOException {
}
}
private void resetIOContext() {
conf.set(Utilities.INPUT_NAME, "TestHiveBinarySearchRecordReader");
ioContext = IOContextMap.get(conf);
ioContext.setUseSorted(false);
ioContext.setBinarySearching(false);
ioContext.setEndBinarySearch(false);
ioContext.setComparison(null);
ioContext.setGenericUDFClassName(null);
}
private void init() throws IOException {
conf = new JobConf();
resetIOContext();
rcfReader = mock(RCFileRecordReader.class);
when(rcfReader.next((LongWritable)anyObject(),
(BytesRefArrayWritable )anyObject())).thenReturn(true);
// Since the start is 0, and the length is 100, the first call to sync should be with the value
// 50 so return that for getPos()
when(rcfReader.getPos()).thenReturn(50L);
conf.setBoolean("hive.input.format.sorted", true);
TableDesc tblDesc = Utilities.defaultTd;
PartitionDesc partDesc = new PartitionDesc(tblDesc, null);
LinkedHashMap<Path, PartitionDesc> pt = new LinkedHashMap<>();
pt.put(new Path("/tmp/testfolder"), partDesc);
MapredWork mrwork = new MapredWork();
mrwork.getMapWork().setPathToPartitionInfo(pt);
Utilities.setMapRedWork(conf, mrwork,new Path("/tmp/" + System.getProperty("user.name"), "hive"));
hiveSplit = new TestHiveInputSplit();
hbsReader = new TestHiveRecordReader(rcfReader, conf);
hbsReader.initIOContext(hiveSplit, conf, Class.class, rcfReader);
}
private boolean executeDoNext(HiveContextAwareRecordReader hbsReader) throws IOException {
return hbsReader.next(hbsReader.createKey(), hbsReader.createValue());
}
public void testNonLinearGreaterThan() throws Exception {
init();
Assert.assertTrue(executeDoNext(hbsReader));
verify(rcfReader).sync(50);
ioContext.setComparison(1);
when(rcfReader.getPos()).thenReturn(25L);
// By setting the comparison to greater, the search should use the block [0, 50]
Assert.assertTrue(executeDoNext(hbsReader));
verify(rcfReader).sync(25);
}
public void testNonLinearLessThan() throws Exception {
init();
Assert.assertTrue(executeDoNext(hbsReader));
verify(rcfReader).sync(50);
ioContext.setComparison(-1);
when(rcfReader.getPos()).thenReturn(75L);
// By setting the comparison to less, the search should use the block [50, 100]
Assert.assertTrue(executeDoNext(hbsReader));
verify(rcfReader).sync(75);
}
public void testNonLinearEqualTo() throws Exception {
init();
Assert.assertTrue(executeDoNext(hbsReader));
verify(rcfReader).sync(50);
ioContext.setComparison(0);
when(rcfReader.getPos()).thenReturn(25L);
// By setting the comparison to equal, the search should use the block [0, 50]
Assert.assertTrue(executeDoNext(hbsReader));
verify(rcfReader).sync(25);
}
public void testHitLastBlock() throws Exception {
init();
Assert.assertTrue(executeDoNext(hbsReader));
verify(rcfReader).sync(50);
ioContext.setComparison(-1);
when(rcfReader.getPos()).thenReturn(100L);
// When sync is called it will return 100, the value signaling the end of the file, this should
// result in a call to sync to the beginning of the block it was searching [50, 100], and it
// should continue normally
Assert.assertTrue(executeDoNext(hbsReader));
InOrder inOrder = inOrder(rcfReader);
inOrder.verify(rcfReader).sync(75);
inOrder.verify(rcfReader).sync(50);
Assert.assertFalse(ioContext.isBinarySearching());
}
public void testHitSamePositionTwice() throws Exception {
init();
Assert.assertTrue(executeDoNext(hbsReader));
verify(rcfReader).sync(50);
ioContext.setComparison(1);
// When getPos is called it should return the same value, signaling the end of the search, so
// the search should continue linearly and it should sync to the beginning of the block [0, 50]
Assert.assertTrue(executeDoNext(hbsReader));
InOrder inOrder = inOrder(rcfReader);
inOrder.verify(rcfReader).sync(25);
inOrder.verify(rcfReader).sync(0);
Assert.assertFalse(ioContext.isBinarySearching());
}
public void testResetRange() throws Exception {
init();
InOrder inOrder = inOrder(rcfReader);
Assert.assertTrue(executeDoNext(hbsReader));
inOrder.verify(rcfReader).sync(50);
ioContext.setComparison(-1);
when(rcfReader.getPos()).thenReturn(75L);
Assert.assertTrue(executeDoNext(hbsReader));
inOrder.verify(rcfReader).sync(75);
ioContext.setEndBinarySearch(true);
// This should make the search linear, sync to the beginning of the block being searched
// [50, 100], set the comparison to be null, and the flag to reset the range should be unset
Assert.assertTrue(executeDoNext(hbsReader));
inOrder.verify(rcfReader).sync(50);
Assert.assertFalse(ioContext.isBinarySearching());
Assert.assertFalse(ioContext.shouldEndBinarySearch());
}
public void testEqualOpClass() throws Exception {
init();
ioContext.setGenericUDFClassName(GenericUDFOPEqual.class.getName());
Assert.assertTrue(ioContext.isBinarySearching());
Assert.assertTrue(executeDoNext(hbsReader));
ioContext.setBinarySearching(false);
ioContext.setComparison(-1);
Assert.assertTrue(executeDoNext(hbsReader));
ioContext.setComparison(0);
Assert.assertTrue(executeDoNext(hbsReader));
ioContext.setComparison(1);
Assert.assertFalse(executeDoNext(hbsReader));
}
public void testLessThanOpClass() throws Exception {
init();
ioContext.setGenericUDFClassName(GenericUDFOPLessThan.class.getName());
Assert.assertTrue(executeDoNext(hbsReader));
Assert.assertFalse(ioContext.isBinarySearching());
ioContext.setComparison(-1);
Assert.assertTrue(executeDoNext(hbsReader));
ioContext.setComparison(0);
Assert.assertFalse(executeDoNext(hbsReader));
ioContext.setComparison(1);
Assert.assertFalse(executeDoNext(hbsReader));
}
public void testLessThanOrEqualOpClass() throws Exception {
init();
ioContext.setGenericUDFClassName(GenericUDFOPEqualOrLessThan.class.getName());
Assert.assertTrue(executeDoNext(hbsReader));
Assert.assertFalse(ioContext.isBinarySearching());
ioContext.setComparison(-1);
Assert.assertTrue(executeDoNext(hbsReader));
ioContext.setComparison(0);
Assert.assertTrue(executeDoNext(hbsReader));
ioContext.setComparison(1);
Assert.assertFalse(executeDoNext(hbsReader));
}
public void testGreaterThanOpClass() throws Exception {
init();
ioContext.setGenericUDFClassName(GenericUDFOPGreaterThan.class.getName());
Assert.assertTrue(ioContext.isBinarySearching());
Assert.assertTrue(executeDoNext(hbsReader));
ioContext.setBinarySearching(false);
ioContext.setComparison(-1);
Assert.assertTrue(executeDoNext(hbsReader));
ioContext.setComparison(0);
Assert.assertTrue(executeDoNext(hbsReader));
ioContext.setComparison(1);
Assert.assertTrue(executeDoNext(hbsReader));
}
public void testGreaterThanOrEqualOpClass() throws Exception {
init();
ioContext.setGenericUDFClassName(GenericUDFOPEqualOrGreaterThan.class.getName());
Assert.assertTrue(ioContext.isBinarySearching());
Assert.assertTrue(executeDoNext(hbsReader));
ioContext.setBinarySearching(false);
ioContext.setComparison(-1);
Assert.assertTrue(executeDoNext(hbsReader));
ioContext.setComparison(0);
Assert.assertTrue(executeDoNext(hbsReader));
ioContext.setComparison(1);
Assert.assertTrue(executeDoNext(hbsReader));
}
public static void main(String[] args) throws Exception {
new TestHiveBinarySearchRecordReader().testNonLinearGreaterThan();
new TestHiveBinarySearchRecordReader().testNonLinearLessThan();
new TestHiveBinarySearchRecordReader().testNonLinearEqualTo();
new TestHiveBinarySearchRecordReader().testHitLastBlock();
new TestHiveBinarySearchRecordReader().testHitSamePositionTwice();
new TestHiveBinarySearchRecordReader().testResetRange();
new TestHiveBinarySearchRecordReader().testEqualOpClass();
new TestHiveBinarySearchRecordReader().testLessThanOpClass();
new TestHiveBinarySearchRecordReader().testLessThanOrEqualOpClass();
new TestHiveBinarySearchRecordReader().testGreaterThanOpClass();
new TestHiveBinarySearchRecordReader().testGreaterThanOrEqualOpClass();
}
}