src/main/java/org/apache/accumulo/testing/randomwalk/bulk/BulkPlusOne.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.accumulo.testing.randomwalk.bulk;

import java.util.ArrayList;
import java.util.List;
import java.util.Random;
import java.util.TreeSet;
import java.util.UUID;
import java.util.concurrent.atomic.AtomicLong;

import org.apache.accumulo.core.client.IteratorSetting.Column;
import org.apache.accumulo.core.client.rfile.RFile;
import org.apache.accumulo.core.client.rfile.RFileWriter;
import org.apache.accumulo.core.data.Key;
import org.apache.accumulo.core.data.Value;
import org.apache.accumulo.testing.randomwalk.RandWalkEnv;
import org.apache.accumulo.testing.randomwalk.State;
import org.apache.hadoop.fs.FileSystem;
import org.apache.hadoop.fs.Path;
import org.apache.hadoop.io.Text;
import org.slf4j.Logger;

public class BulkPlusOne extends BulkImportTest {

  public static final int LOTS = 100000;
  public static final int COLS = 10;
  public static final int HEX_SIZE = (int) Math.ceil(Math.log(LOTS) / Math.log(16));
  public static final String FMT = "r%0" + HEX_SIZE + "x";
  public static final List<Column> COLNAMES = new ArrayList<>();
  public static final Text CHECK_COLUMN_FAMILY = new Text("cf");
  static {
    for (int i = 0; i < COLS; i++) {
      COLNAMES.add(new Column(CHECK_COLUMN_FAMILY, new Text(String.format("%03d", i))));
    }
  }
  public static final Text MARKER_CF = new Text("marker");
  static final AtomicLong counter = new AtomicLong();

  private static final Value ONE = new Value("1".getBytes());

  static void bulkLoadLots(Logger log, State state, RandWalkEnv env, Value value) throws Exception {
    final FileSystem fs = (FileSystem) state.get("fs");
    final Path dir = new Path(fs.getUri() + "/tmp", "bulk_" + UUID.randomUUID().toString());
    log.debug("Bulk loading from {}", dir);
    final Random rand = state.getRandom();
    final int parts = rand.nextInt(10) + 1;

    TreeSet<Integer> startRows = new TreeSet<>();
    startRows.add(0);
    while (startRows.size() < parts)
      startRows.add(rand.nextInt(LOTS));

    List<String> printRows = new ArrayList<>(startRows.size());
    for (Integer row : startRows)
      printRows.add(String.format(FMT, row));

    String markerColumnQualifier = String.format("%07d", counter.incrementAndGet());
    log.debug("preparing bulk files with start rows " + printRows + " last row "
        + String.format(FMT, LOTS - 1) + " marker " + markerColumnQualifier);

    List<Integer> rows = new ArrayList<>(startRows);
    rows.add(LOTS);

    for (int i = 0; i < parts; i++) {
      String fileName = dir + "/" + String.format("part_%d.rf", i);

      log.debug("Creating {}", fileName);
      RFileWriter writer = RFile.newWriter().to(fileName).withFileSystem(fs).build();
      writer.startDefaultLocalityGroup();
      int start = rows.get(i);
      int end = rows.get(i + 1);
      for (int j = start; j < end; j++) {
        Text row = new Text(String.format(FMT, j));
        for (Column col : COLNAMES) {
          writer.append(new Key(row, col.getColumnFamily(), col.getColumnQualifier()), value);
        }
        writer.append(new Key(row, MARKER_CF, new Text(markerColumnQualifier)), ONE);
      }
      writer.close();
    }
    env.getAccumuloClient().tableOperations().importDirectory(dir.toString())
        .to(Setup.getTableName()).tableTime(true);
    fs.delete(dir, true);
    log.debug("Finished bulk import, start rows " + printRows + " last row "
        + String.format(FMT, LOTS - 1) + " marker " + markerColumnQualifier);
  }

  @Override
  protected void runLater(State state, RandWalkEnv env) throws Exception {
    log.info("Incrementing");
    bulkLoadLots(log, state, env, ONE);
  }

}