tajo-core/tajo-core-backend/src/test/java/org/apache/tajo/worker/TestRangeRetrieverHandler.java

/**
 * 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.tajo.worker;

import com.google.common.collect.Lists;
import com.google.common.collect.Maps;
import org.apache.commons.codec.binary.Base64;
import org.apache.hadoop.fs.FileSystem;
import org.apache.hadoop.fs.Path;
import org.apache.tajo.LocalTajoTestingUtility;
import org.apache.tajo.TajoTestingCluster;
import org.apache.tajo.algebra.Expr;
import org.apache.tajo.catalog.*;
import org.apache.tajo.catalog.proto.CatalogProtos.StoreType;
import org.apache.tajo.common.TajoDataTypes.Type;
import org.apache.tajo.conf.TajoConf;
import org.apache.tajo.datum.Datum;
import org.apache.tajo.datum.DatumFactory;
import org.apache.tajo.engine.parser.SQLAnalyzer;
import org.apache.tajo.engine.planner.*;
import org.apache.tajo.engine.planner.enforce.Enforcer;
import org.apache.tajo.engine.planner.logical.LogicalNode;
import org.apache.tajo.engine.planner.physical.ExternalSortExec;
import org.apache.tajo.engine.planner.physical.PhysicalExec;
import org.apache.tajo.engine.planner.physical.ProjectionExec;
import org.apache.tajo.engine.planner.physical.RangeShuffleFileWriteExec;
import org.apache.tajo.storage.*;
import org.apache.tajo.storage.fragment.FileFragment;
import org.apache.tajo.storage.index.bst.BSTIndex;
import org.apache.tajo.util.CommonTestingUtil;
import org.apache.tajo.worker.dataserver.retriever.FileChunk;
import org.junit.After;
import org.junit.Before;
import org.junit.Test;

import java.io.File;
import java.io.IOException;
import java.util.ArrayList;
import java.util.List;
import java.util.Map;

import static org.junit.Assert.assertEquals;
import static org.junit.Assert.assertTrue;

public class TestRangeRetrieverHandler {
  private TajoTestingCluster util;
  private TajoConf conf;
  private CatalogService catalog;
  private SQLAnalyzer analyzer;
  private LogicalPlanner planner;
  private LogicalOptimizer optimizer;
  private AbstractStorageManager sm;
  private Schema schema;
  private static int TEST_TUPLE = 10000;
  private FileSystem fs;
  private Path testDir;

  @Before
  public void setUp() throws Exception {
    util = new TajoTestingCluster();
    conf = util.getConfiguration();
    testDir = CommonTestingUtil.getTestDir("target/test-data/TestRangeRetrieverHandler");
    fs = testDir.getFileSystem(conf);
    util.startCatalogCluster();
    catalog = util.getMiniCatalogCluster().getCatalog();
    sm = StorageManagerFactory.getStorageManager(conf, testDir);

    analyzer = new SQLAnalyzer();
    planner = new LogicalPlanner(catalog);
    optimizer = new LogicalOptimizer(conf);

    schema = new Schema();
    schema.addColumn("empId", Type.INT4);
    schema.addColumn("age", Type.INT4);
  }

  @After
  public void tearDown() {
    util.shutdownCatalogCluster();
  }

  public String [] SORT_QUERY = {
      "select empId, age from employee order by empId, age",
      "select empId, age from employee order by empId desc, age desc"
  };

  @Test
  public void testGet() throws Exception {
    Tuple firstTuple = null;
    Tuple lastTuple;

    TableMeta employeeMeta = CatalogUtil.newTableMeta(StoreType.CSV);

    Path tableDir = StorageUtil.concatPath(testDir, "testGet", "table.csv");
    fs.mkdirs(tableDir.getParent());
    Appender appender = sm.getAppender(employeeMeta, schema, tableDir);
    appender.init();

    Tuple tuple = new VTuple(schema.size());
    for (int i = 0; i < TEST_TUPLE; i++) {
      tuple.put(
          new Datum[] {
              DatumFactory.createInt4(i),
              DatumFactory.createInt4(i + 5)
          });
      appender.addTuple(tuple);

      if (firstTuple == null) {
        firstTuple = new VTuple(tuple);
      }
    }
    lastTuple = new VTuple(tuple);
    appender.flush();
    appender.close();

    TableDesc employee = new TableDesc("employee", schema, employeeMeta, tableDir);
    catalog.addTable(employee);

    FileFragment[] frags = StorageManager.splitNG(conf, "employee", employeeMeta, tableDir, Integer.MAX_VALUE);

    TaskAttemptContext ctx = new TaskAttemptContext(conf, LocalTajoTestingUtility.newQueryUnitAttemptId(),
        new FileFragment[] {frags[0]}, testDir);
    ctx.setEnforcer(new Enforcer());
    Expr expr = analyzer.parse(SORT_QUERY[0]);
    LogicalPlan plan = planner.createPlan(expr);
    LogicalNode rootNode = optimizer.optimize(plan);

    PhysicalPlanner phyPlanner = new PhysicalPlannerImpl(conf,sm);
    PhysicalExec exec = phyPlanner.createPlan(ctx, rootNode);

    ExternalSortExec sort = null;
    if (exec instanceof ProjectionExec) {
      ProjectionExec projExec = (ProjectionExec) exec;
      sort = projExec.getChild();
    } else if (exec instanceof ExternalSortExec) {
      sort = (ExternalSortExec) exec;
    } else {
      assertTrue(false);
    }

    SortSpec[] sortSpecs = sort.getPlan().getSortKeys();
    RangeShuffleFileWriteExec idxStoreExec = new RangeShuffleFileWriteExec(ctx, sm, sort, sort.getSchema(),
        sort.getSchema(), sortSpecs);

    exec = idxStoreExec;
    exec.init();
    exec.next();
    exec.close();

    Schema keySchema = PlannerUtil.sortSpecsToSchema(sortSpecs);
    TupleComparator comp = new TupleComparator(keySchema, sortSpecs);
    BSTIndex bst = new BSTIndex(conf);
    BSTIndex.BSTIndexReader reader = bst.getIndexReader(
        new Path(testDir, "output/index"), keySchema, comp);
    reader.open();

    TableMeta meta = CatalogUtil.newTableMeta(StoreType.RAW, new Options());
    SeekableScanner scanner = StorageManagerFactory.getSeekableScanner(conf, meta, schema,
        StorageUtil.concatPath(testDir, "output", "output"));

    scanner.init();
    int cnt = 0;
    while(scanner.next() != null) {
      cnt++;
    }
    scanner.reset();

    assertEquals(TEST_TUPLE ,cnt);

    Tuple keytuple = new VTuple(2);
    for(int i = 1 ; i < TEST_TUPLE ; i ++) {
      keytuple.put(0, DatumFactory.createInt4(i));
      keytuple.put(1, DatumFactory.createInt4(i + 5));
      long offsets = reader.find(keytuple);
      scanner.seek(offsets);
      tuple = scanner.next();
      assertTrue("[seek check " + (i) + " ]" , i == tuple.get(0).asInt4());
      //assertTrue("[seek check " + (i) + " ]" , ("name_" + i).equals(tuple.get(1).asChars()));
    }

    TupleRange totalRange = new TupleRange(sortSpecs, firstTuple, lastTuple);
    UniformRangePartition partitioner = new UniformRangePartition(totalRange, sortSpecs, true);
    TupleRange [] partitions = partitioner.partition(7);

    // The below is for testing RangeRetrieverHandler.
    RangeRetrieverHandler handler = new RangeRetrieverHandler(
        new File((new Path(testDir, "output")).toUri()), keySchema, comp);

    List<Long []> offsets = new ArrayList<Long []>();

    for (int i = 0; i < partitions.length; i++) {
      FileChunk chunk = getFileChunk(handler, keySchema, partitions[i], i == (partitions.length - 1));
      offsets.add(new Long[] {chunk.startOffset(), chunk.length()});
    }
    scanner.close();

    Long[] previous = null;
    for (Long [] offset : offsets) {
      if (offset[0] == 0 && previous == null) {
        previous = offset;
        continue;
      }
      assertTrue(previous[0] + previous[1] == offset[0]);
      previous = offset;
    }
    long fileLength = new File((new Path(testDir, "index").toUri())).length();
    assertTrue(previous[0] + previous[1] == fileLength);
  }

  @Test
  public void testGetFromDescendingOrder() throws Exception {
    Tuple firstTuple = null;
    Tuple lastTuple;

    TableMeta meta = CatalogUtil.newTableMeta(StoreType.CSV);
    Path tablePath = StorageUtil.concatPath(testDir, "testGetFromDescendingOrder", "table.csv");
    fs.mkdirs(tablePath.getParent());
    Appender appender = sm.getAppender(meta, schema, tablePath);
    appender.init();
    Tuple tuple = new VTuple(schema.size());
    for (int i = (TEST_TUPLE - 1); i >= 0 ; i--) {
      tuple.put(
          new Datum[] {
              DatumFactory.createInt4(i),
              DatumFactory.createInt4(i + 5)
          });
      appender.addTuple(tuple);

      if (firstTuple == null) {
        firstTuple = new VTuple(tuple);
      }
    }
    lastTuple = new VTuple(tuple);
    appender.flush();
    appender.close();

    TableDesc employee = new TableDesc("employee", schema, meta, tablePath);
    catalog.addTable(employee);

    FileFragment[] frags = sm.splitNG(conf, "employee", meta, tablePath, Integer.MAX_VALUE);

    TaskAttemptContext
        ctx = new TaskAttemptContext(conf, LocalTajoTestingUtility.newQueryUnitAttemptId(),
        new FileFragment[] {frags[0]}, testDir);
    ctx.setEnforcer(new Enforcer());
    Expr expr = analyzer.parse(SORT_QUERY[1]);
    LogicalPlan plan = planner.createPlan(expr);
    LogicalNode rootNode = optimizer.optimize(plan);

    PhysicalPlanner phyPlanner = new PhysicalPlannerImpl(conf,sm);
    PhysicalExec exec = phyPlanner.createPlan(ctx, rootNode);

    ExternalSortExec sort = null;
    if (exec instanceof ProjectionExec) {
      ProjectionExec projExec = (ProjectionExec) exec;
      sort = projExec.getChild();
    } else if (exec instanceof ExternalSortExec) {
      sort = (ExternalSortExec) exec;
    } else {
      assertTrue(false);
    }

    SortSpec[] sortSpecs = sort.getPlan().getSortKeys();
    RangeShuffleFileWriteExec idxStoreExec = new RangeShuffleFileWriteExec(ctx, sm, sort,
        sort.getSchema(), sort.getSchema(), sortSpecs);

    exec = idxStoreExec;
    exec.init();
    exec.next();
    exec.close();

    Schema keySchema = PlannerUtil.sortSpecsToSchema(sortSpecs);
    TupleComparator comp = new TupleComparator(keySchema, sortSpecs);
    BSTIndex bst = new BSTIndex(conf);
    BSTIndex.BSTIndexReader reader = bst.getIndexReader(
        new Path(testDir, "output/index"), keySchema, comp);
    reader.open();
    TableMeta outputMeta = CatalogUtil.newTableMeta(StoreType.RAW, new Options());
    SeekableScanner scanner = StorageManagerFactory.getSeekableScanner(conf, outputMeta, schema,
        StorageUtil.concatPath(testDir, "output", "output"));
    scanner.init();
    int cnt = 0;
    while(scanner.next() != null) {
      cnt++;
    }
    scanner.reset();

    assertEquals(TEST_TUPLE ,cnt);

    Tuple keytuple = new VTuple(2);
    for(int i = (TEST_TUPLE - 1) ; i >= 0; i --) {
      keytuple.put(0, DatumFactory.createInt4(i));
      keytuple.put(1, DatumFactory.createInt4(i + 5));
      long offsets = reader.find(keytuple);
      scanner.seek(offsets);
      tuple = scanner.next();
      assertTrue("[seek check " + (i) + " ]" , i == tuple.get(0).asInt4());
    }

    TupleRange totalRange = new TupleRange(sortSpecs, firstTuple, lastTuple);
    UniformRangePartition partitioner = new UniformRangePartition(totalRange, sortSpecs, true);
    TupleRange [] partitions = partitioner.partition(25);

    File dataFile = new File((new Path(testDir, "output")).toUri());

    // The below is for testing RangeRetrieverHandler.
    RangeRetrieverHandler handler = new RangeRetrieverHandler(
        dataFile, keySchema, comp);

    List<Long []> offsets = new ArrayList<Long []>();

    for (int i = 0; i < partitions.length; i++) {
      FileChunk chunk = getFileChunk(handler, keySchema, partitions[i], i == 0);
      offsets.add(new Long[] {chunk.startOffset(), chunk.length()});
    }
    scanner.close();

    long fileLength = new File(dataFile, "data/data").length();
    Long[] previous = null;
    for (Long [] offset : offsets) {
      if (previous == null) {
        assertTrue(offset[0] + offset[1] == fileLength);
        previous = offset;
        continue;
      }

      assertTrue(offset[0] + offset[1] == previous[0]);
      previous = offset;
    }
  }

  private FileChunk getFileChunk(RangeRetrieverHandler handler, Schema keySchema,
                                 TupleRange range, boolean last) throws IOException {
    Map<String,List<String>> kvs = Maps.newHashMap();
    kvs.put("start", Lists.newArrayList(
        new String(Base64.encodeBase64(
            RowStoreUtil.RowStoreEncoder.toBytes(keySchema, range.getStart()),
            false))));
    kvs.put("end", Lists.newArrayList(
        new String(Base64.encodeBase64(
            RowStoreUtil.RowStoreEncoder.toBytes(keySchema, range.getEnd()), false))));

    if (last) {
      kvs.put("final", Lists.newArrayList("true"));
    }
    return handler.get(kvs);
  }
}