hyracks-examples/hyracks-integration-tests/src/test/java/edu/uci/ics/hyracks/tests/btree/BTreePrimaryIndexStatsOperatorTest.java

/*
 * Copyright 2009-2010 by The Regents of the University of California
 * 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 from
 * 
 *     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 edu.uci.ics.hyracks.tests.btree;

import java.io.File;
import java.text.SimpleDateFormat;
import java.util.Date;

import org.junit.AfterClass;
import org.junit.Before;
import org.junit.Test;

import edu.uci.ics.hyracks.api.constraints.PartitionConstraintHelper;
import edu.uci.ics.hyracks.api.dataflow.value.IBinaryComparatorFactory;
import edu.uci.ics.hyracks.api.dataflow.value.ISerializerDeserializer;
import edu.uci.ics.hyracks.api.dataflow.value.ITypeTrait;
import edu.uci.ics.hyracks.api.dataflow.value.RecordDescriptor;
import edu.uci.ics.hyracks.api.dataflow.value.TypeTrait;
import edu.uci.ics.hyracks.api.io.FileReference;
import edu.uci.ics.hyracks.api.job.JobSpecification;
import edu.uci.ics.hyracks.dataflow.common.data.comparators.UTF8StringBinaryComparatorFactory;
import edu.uci.ics.hyracks.dataflow.common.data.marshalling.UTF8StringSerializerDeserializer;
import edu.uci.ics.hyracks.dataflow.common.data.parsers.IValueParserFactory;
import edu.uci.ics.hyracks.dataflow.common.data.parsers.UTF8StringParserFactory;
import edu.uci.ics.hyracks.dataflow.std.connectors.OneToOneConnectorDescriptor;
import edu.uci.ics.hyracks.dataflow.std.file.ConstantFileSplitProvider;
import edu.uci.ics.hyracks.dataflow.std.file.DelimitedDataTupleParserFactory;
import edu.uci.ics.hyracks.dataflow.std.file.FileScanOperatorDescriptor;
import edu.uci.ics.hyracks.dataflow.std.file.FileSplit;
import edu.uci.ics.hyracks.dataflow.std.file.IFileSplitProvider;
import edu.uci.ics.hyracks.dataflow.std.sort.ExternalSortOperatorDescriptor;
import edu.uci.ics.hyracks.storage.am.btree.dataflow.BTreeOpHelperFactory;
import edu.uci.ics.hyracks.storage.am.btree.frames.BTreeNSMInteriorFrameFactory;
import edu.uci.ics.hyracks.storage.am.btree.frames.BTreeNSMLeafFrameFactory;
import edu.uci.ics.hyracks.storage.am.common.api.ITreeIndex;
import edu.uci.ics.hyracks.storage.am.common.api.ITreeIndexFrameFactory;
import edu.uci.ics.hyracks.storage.am.common.dataflow.IIndexRegistryProvider;
import edu.uci.ics.hyracks.storage.am.common.dataflow.ITreeIndexOpHelperFactory;
import edu.uci.ics.hyracks.storage.am.common.dataflow.TreeIndexBulkLoadOperatorDescriptor;
import edu.uci.ics.hyracks.storage.am.common.dataflow.TreeIndexStatsOperatorDescriptor;
import edu.uci.ics.hyracks.storage.am.common.tuples.TypeAwareTupleWriterFactory;
import edu.uci.ics.hyracks.storage.common.IStorageManagerInterface;
import edu.uci.ics.hyracks.test.support.TestStorageManagerComponentHolder;
import edu.uci.ics.hyracks.test.support.TestStorageManagerInterface;
import edu.uci.ics.hyracks.test.support.TestTreeIndexRegistryProvider;
import edu.uci.ics.hyracks.tests.integration.AbstractIntegrationTest;

public class BTreePrimaryIndexStatsOperatorTest extends AbstractIntegrationTest {
	static {
		TestStorageManagerComponentHolder.init(8192, 20, 20);
	}

	private IStorageManagerInterface storageManager = new TestStorageManagerInterface();
	private IIndexRegistryProvider<ITreeIndex> treeIndexRegistryProvider = new TestTreeIndexRegistryProvider();
	private ITreeIndexOpHelperFactory opHelperFactory = new BTreeOpHelperFactory();

	private final static SimpleDateFormat simpleDateFormat = new SimpleDateFormat(
			"ddMMyy-hhmmssSS");
	private final static String sep = System.getProperty("file.separator");

	// field, type and key declarations for primary index
	private int primaryFieldCount = 6;
	private ITypeTrait[] primaryTypeTraits = new ITypeTrait[primaryFieldCount];
	private int primaryKeyFieldCount = 1;
	private IBinaryComparatorFactory[] primaryComparatorFactories = new IBinaryComparatorFactory[primaryKeyFieldCount];
	private TypeAwareTupleWriterFactory primaryTupleWriterFactory = new TypeAwareTupleWriterFactory(
			primaryTypeTraits);
	private ITreeIndexFrameFactory primaryInteriorFrameFactory = new BTreeNSMInteriorFrameFactory(
			primaryTupleWriterFactory);
	private ITreeIndexFrameFactory primaryLeafFrameFactory = new BTreeNSMLeafFrameFactory(
			primaryTupleWriterFactory);

	private static String primaryBtreeName = "primary"
			+ simpleDateFormat.format(new Date());
	private static String primaryFileName = System
			.getProperty("java.io.tmpdir") + sep + primaryBtreeName;

	private IFileSplitProvider primaryBtreeSplitProvider = new ConstantFileSplitProvider(
			new FileSplit[] { new FileSplit(NC1_ID, new FileReference(new File(
					primaryFileName))) });

	private RecordDescriptor primaryRecDesc = new RecordDescriptor(
			new ISerializerDeserializer[] {
					UTF8StringSerializerDeserializer.INSTANCE,
					UTF8StringSerializerDeserializer.INSTANCE,
					UTF8StringSerializerDeserializer.INSTANCE,
					UTF8StringSerializerDeserializer.INSTANCE,
					UTF8StringSerializerDeserializer.INSTANCE,
					UTF8StringSerializerDeserializer.INSTANCE });

	@Before
	public void setup() throws Exception {
		// field, type and key declarations for primary index
		primaryTypeTraits[0] = new TypeTrait(ITypeTrait.VARIABLE_LENGTH);
		primaryTypeTraits[1] = new TypeTrait(ITypeTrait.VARIABLE_LENGTH);
		primaryTypeTraits[2] = new TypeTrait(ITypeTrait.VARIABLE_LENGTH);
		primaryTypeTraits[3] = new TypeTrait(ITypeTrait.VARIABLE_LENGTH);
		primaryTypeTraits[4] = new TypeTrait(ITypeTrait.VARIABLE_LENGTH);
		primaryTypeTraits[5] = new TypeTrait(ITypeTrait.VARIABLE_LENGTH);
		primaryComparatorFactories[0] = UTF8StringBinaryComparatorFactory.INSTANCE;

		loadPrimaryIndexTest();
	}

	public void loadPrimaryIndexTest() throws Exception {
		JobSpecification spec = new JobSpecification();

		FileSplit[] ordersSplits = new FileSplit[] { new FileSplit(NC1_ID,
				new FileReference(new File("data/tpch0.001/orders-part1.tbl"))) };
		IFileSplitProvider ordersSplitProvider = new ConstantFileSplitProvider(
				ordersSplits);
		RecordDescriptor ordersDesc = new RecordDescriptor(
				new ISerializerDeserializer[] {
						UTF8StringSerializerDeserializer.INSTANCE,
						UTF8StringSerializerDeserializer.INSTANCE,
						UTF8StringSerializerDeserializer.INSTANCE,
						UTF8StringSerializerDeserializer.INSTANCE,
						UTF8StringSerializerDeserializer.INSTANCE,
						UTF8StringSerializerDeserializer.INSTANCE,
						UTF8StringSerializerDeserializer.INSTANCE,
						UTF8StringSerializerDeserializer.INSTANCE,
						UTF8StringSerializerDeserializer.INSTANCE });

		FileScanOperatorDescriptor ordScanner = new FileScanOperatorDescriptor(
				spec, ordersSplitProvider, new DelimitedDataTupleParserFactory(
						new IValueParserFactory[] {
								UTF8StringParserFactory.INSTANCE,
								UTF8StringParserFactory.INSTANCE,
								UTF8StringParserFactory.INSTANCE,
								UTF8StringParserFactory.INSTANCE,
								UTF8StringParserFactory.INSTANCE,
								UTF8StringParserFactory.INSTANCE,
								UTF8StringParserFactory.INSTANCE,
								UTF8StringParserFactory.INSTANCE,
								UTF8StringParserFactory.INSTANCE }, '|'),
				ordersDesc);
		PartitionConstraintHelper.addAbsoluteLocationConstraint(spec,
				ordScanner, NC1_ID);

		ExternalSortOperatorDescriptor sorter = new ExternalSortOperatorDescriptor(
				spec,
				1000,
				new int[] { 0 },
				new IBinaryComparatorFactory[] { UTF8StringBinaryComparatorFactory.INSTANCE },
				ordersDesc);
		PartitionConstraintHelper.addAbsoluteLocationConstraint(spec, sorter,
				NC1_ID);

		int[] fieldPermutation = { 0, 1, 2, 4, 5, 7 };
		TreeIndexBulkLoadOperatorDescriptor primaryBtreeBulkLoad = new TreeIndexBulkLoadOperatorDescriptor(
				spec, storageManager, treeIndexRegistryProvider,
				primaryBtreeSplitProvider, primaryInteriorFrameFactory,
				primaryLeafFrameFactory, primaryTypeTraits,
				primaryComparatorFactories, null, fieldPermutation, 0.7f,
				opHelperFactory);
		PartitionConstraintHelper.addAbsoluteLocationConstraint(spec,
				primaryBtreeBulkLoad, NC1_ID);

		spec.connect(new OneToOneConnectorDescriptor(spec), ordScanner, 0,
				sorter, 0);

		spec.connect(new OneToOneConnectorDescriptor(spec), sorter, 0,
				primaryBtreeBulkLoad, 0);

		spec.addRoot(primaryBtreeBulkLoad);
		runTest(spec);
	}

	@Test
	public void showPrimaryIndexStats() throws Exception {
		JobSpecification spec = new JobSpecification();

		TreeIndexStatsOperatorDescriptor primaryStatsOp = new TreeIndexStatsOperatorDescriptor(
				spec, storageManager, treeIndexRegistryProvider,
				primaryBtreeSplitProvider, primaryInteriorFrameFactory,
				primaryLeafFrameFactory, primaryTypeTraits,
				primaryComparatorFactories, opHelperFactory);
		PartitionConstraintHelper.addAbsoluteLocationConstraint(spec,
				primaryStatsOp, NC1_ID);

		spec.addRoot(primaryStatsOp);
		runTest(spec);
	}

	@AfterClass
	public static void cleanup() throws Exception {
		File primary = new File(primaryFileName);
		primary.deleteOnExit();
	}
}