org.apache.cassandra.io.sstable.CQLSSTableWriter Java Examples

/**
 * The column values must correspond to the order in which
 * they appear in the insert stored procedure. 
 * 
 * Key is not used, so it can be null or any object.
 * </p>
 *
 * @param key
 *            any object or null.
 * @param values
 *            the values to write.
 * @throws IOException
 */
@Override
public void write(Object key, List<ByteBuffer> values) throws IOException
{
    prepareWriter();
    try
    {
        ((CQLSSTableWriter) writer).rawAddRow(values);
        
        if (null != progress)
            progress.progress();
        if (null != context)
            HadoopCompat.progress(context);
    } 
    catch (InvalidRequestException e)
    {
        throw new IOException("Error adding row with key: " + key, e);
    }
}
 

/**
 * Builds an SSTableWriter for a static table
 * @param sorted Defines if the SSTableWriters created by this should be set as sorted, which improves performance if
 *               rows are inserted in SSTable sort order, but throws exceptions if they are inserted in the wrong order.
 * @return A CQLSSTableWriter for this static table
 * @throws CQLGenerationException
 * @throws IOException
 */
private CQLSSTableWriter buildSSTableWriterForShardIndexTable(boolean sorted) throws CQLGenerationException, IOException {
	// Generate CQL create syntax
	String createCQL = this.cqlGenerator.makeCQLforShardIndexTableCreate().getQuery();

	// Generate CQL insert syntax
	String tableName = CObjectShardList.SHARD_INDEX_TABLE_NAME;
	String insertCQL = this.cqlGenerator.makeCQLforInsertNoValuesforShardIndex(tableName).getQuery();

	String SSTablePath = this.defaultSSTableOutputPath + "/" + keyspaceDefinition.getName() + "/" + tableName;
	if (!new File(SSTablePath).mkdir()) {
		throw new IOException("Failed to create new directory for SSTable writing at path: " + SSTablePath);
	}

	CQLSSTableWriter.Builder builder =
			CQLSSTableWriter.builder()
					.inDirectory(SSTablePath)
					.forTable(createCQL)
					.using(insertCQL);
	if (sorted) { builder = builder.sorted(); }
	return builder.build();
}
 

/**
 * Builds an SSTableWriter for a static table
 * @param definition Definition of object to build table for
 * @param sorted Defines if the SSTableWriters created by this should be set as sorted, which improves performance if
 *               rows are inserted in SSTable sort order, but throws exceptions if they are inserted in the wrong order.
 * @return A CQLSSTableWriter for this static table
 * @throws CQLGenerationException
 * @throws IOException
 */
private CQLSSTableWriter buildSSTableWriterForStaticTable(CDefinition definition, boolean sorted) throws CQLGenerationException, IOException {
	// Generate CQL create syntax
	String tableName = definition.getName();
	String createCQL = this.cqlGenerator.makeStaticTableCreate(definition).getQuery();

	// Generate CQL insert syntax
	String insertCQL = this.cqlGenerator.makeCQLforInsertNoValuesforStaticTable(tableName).getQuery();

	String SSTablePath = this.defaultSSTableOutputPath + "/" + keyspaceDefinition.getName() + "/" + tableName;
	if (!new File(SSTablePath).mkdir()) {
		throw new IOException("Failed to create new directory for SSTable writing at path: " + SSTablePath);
	}

	CQLSSTableWriter.Builder builder =
		CQLSSTableWriter.builder()
			.inDirectory(SSTablePath)
			.forTable(createCQL)
			.using(insertCQL);
	if (sorted) { builder = builder.sorted(); }
	return builder.build();
}
 

/**
 * Builds an SSTableWriter for a wide/index table
 * @param definition The definition this index/wide table is on
 * @param index The index/wide table to create an CQLSSTableWriter for
 * @param sorted Defines if the SSTableWriters created by this should be set as sorted, which improves performance if
 *               rows are inserted in SSTable sort order, but throws exceptions if they are inserted in the wrong order.
 * @return An CQLSSTableWriter for this wide table
 * @throws CQLGenerationException
 */
private CQLSSTableWriter buildSSTableWriterForWideTable(CDefinition definition, CIndex index, boolean sorted) throws CQLGenerationException, IOException {
	String indexTableName = CObjectCQLGenerator.makeTableName(definition, index);
	// Generate CQL create syntax
	String createCQL = this.cqlGenerator.makeWideTableCreate(definition, index).getQuery();

	// Generate CQL insert syntax
	// Just use 1 as the value for shardId, doesn't matter since we're not actually using values here
	String insertCQL = this.cqlGenerator.makeCQLforInsertNoValuesforWideTable(definition, indexTableName, 1L).getQuery();

	String SSTablePath = this.defaultSSTableOutputPath + "/" + keyspaceDefinition.getName() + "/" + indexTableName;
	if (!new File(SSTablePath).mkdir()) {
		throw new IOException("Failed to create new directory for SSTable writing at path: " + SSTablePath);
	}

	CQLSSTableWriter.Builder builder =
			CQLSSTableWriter.builder()
					.inDirectory(SSTablePath)
					.forTable(createCQL)
					.using(insertCQL);
	if (sorted) { builder = builder.sorted(); }
	return builder.build();
}
 

/**
 * Completes writes to SSTables and cleans up empty table directories. Must be called after writing to SSTables if
 * you actually want to use the SSTables for anything.
 * @throws IOException
 */
public void completeSSTableWrites() throws IOException {
	Map<String, CDefinition> definitions = this.keyspaceDefinition.getDefinitions();
	for (String tableName : this.SSTableWriters.keySet()) {
		// Close all SSTableWriters
		this.SSTableWriters.get(tableName).left.close();
		this.clearSSTableDirectoryIfEmpty(tableName);

		Map<CIndex, CQLSSTableWriter> indexWriters = this.SSTableWriters.get(tableName).right;
		if (indexWriters != null) {
			for (CQLSSTableWriter indexWriter : indexWriters.values()) {
				indexWriter.close();
			}
			// Clear out empty SSTable directories that haven't been written to
			CDefinition def = definitions.get(tableName);
			List<CIndex> indexes = def.getIndexesAsList();
			for (CIndex index : indexes) {
				this.clearSSTableDirectoryIfEmpty(CObjectCQLGenerator.makeTableName(def, index));
			}
		}
	}
}
 

private void prepareWriter()  {
  try {
    if (writer == null) {
      writer = CQLSSTableWriter.builder()
          .forTable(schema)
          .using(insertStatement)
          .withPartitioner(ConfigHelper.getOutputPartitioner(conf))
          .inDirectory(outputDir)
          .sorted()
          .build();
    }
    if (loader == null) {
      CrunchExternalClient externalClient = new CrunchExternalClient(conf);
      externalClient.addKnownCfs(keyspace, schema);
      this.loader = new SSTableLoader(outputDir, externalClient,
          new BulkRecordWriter.NullOutputHandler());
    }
  } catch (Exception e) {
    throw new CrunchRuntimeException(e);
  }
}
 

@Override
public void write(final ByteBuffer ignoredKey, final CQLRecord record)  {
  prepareWriter();
  // To ensure Crunch doesn't reuse CQLSSTableWriter's objects
  List<ByteBuffer> bb = Lists.newArrayList();
  for (ByteBuffer v : record.getValues()) {
    bb.add(ByteBufferUtil.clone(v));
  }
  try {
    ((CQLSSTableWriter) writer).rawAddRow(bb);
    if (null != progress)
      progress.progress();
    if (null != context)
      HadoopCompat.progress(context);
  } catch (InvalidRequestException | IOException e) {
    LOG.error(e.getMessage());
    throw new CrunchRuntimeException("Error adding row : " + e.getMessage());
  }
}
 

private CQLSSTableWriter createWriter() {
        String schema = "CREATE TABLE " + keyspace +  ".data ( " +
                "tenant_id text, " +
                "type tinyint, " +
                "metric text, " +
                "dpart bigint, " +
                "time timeuuid, " +
                "data_retention int static, " +
                "n_value double, " +
                "availability blob, " +
                "l_value bigint, " +
                // Commenting out the aggregates column because there appears to be a bug in the C*
                // code that breaks CQLSSTableWriter when the schema includes a collection. Fortunately,
                // we are not using the column so it can safely be ignored.
//                "aggregates set<frozen <aggregate_data>>, " +
                "PRIMARY KEY ((tenant_id, type, metric, dpart), time) " +
                ") WITH CLUSTERING ORDER BY (time DESC)";

        String insertGauge = "INSERT INTO " + keyspace + ".data (tenant_id, type, metric, dpart, time, n_value) " +
                "VALUES (?, ?, ?, ?, ?, ?)";


        return CQLSSTableWriter.builder()
                .inDirectory(dataDir)
                .forTable(schema)
                .using(insertGauge)
                .withBufferSizeInMB(bufferSize)
                .build();
    }
 

private void prepareWriter() throws IOException
{
    try
    {
        if (writer == null)
        {
            writer = CQLSSTableWriter.builder()
                .forTable(schema)
                .using(insertStatement)
                .withPartitioner(ConfigHelper.getOutputPartitioner(conf))
                .inDirectory(outputDir)
                .withBufferSizeInMB(Integer.parseInt(conf.get(BUFFER_SIZE_IN_MB, "64")))
                .build();
        }
        if (loader == null)
        {
            ExternalClient externalClient = new ExternalClient(conf);
            
            externalClient.addKnownCfs(keyspace, schema);

            this.loader = new SSTableLoader(outputDir, externalClient, new BulkRecordWriter.NullOutputHandler()) {
                @Override
                public void onSuccess(StreamState finalState)
                {
                    if (deleteSrc)
                        FileUtils.deleteRecursive(outputDir);
                }
            };
        }
    }
    catch (Exception e)
    {
        throw new IOException(e);
    }      
}
 

/**
 * Creates an SSTable keyspace output directory at defaultSSTableOutputPath and table output directories for each SSTable,
 * and initializes each SSTableWriter for each static and index table in this keyspace.
 * @param sorted Defines if the SSTableWriters created by this should be set as sorted, which improves performance if
 *               rows are inserted in SSTable sort order, but throws exceptions if they are inserted in the wrong order.
 * @throws CQLGenerationException
 * @throws IOException
 */
public void initializeSSTableWriters(boolean sorted) throws CQLGenerationException, IOException {
	Map<String, CDefinition> definitions = this.keyspaceDefinition.getDefinitions();

	// Make sure the SSTableOutput directory exists and is clear
	String keyspacePath = this.defaultSSTableOutputPath + "/" + this.keyspaceDefinition.getName();
	File keyspaceDir = new File(keyspacePath);
	if (keyspaceDir.exists()) {
		FileUtils.deleteRecursive(new File(keyspacePath));
	}
	if (!new File(keyspacePath).mkdir()) {
		throw new IOException("Failed to create SSTable keyspace output directory at " + keyspacePath);
	}

	this.SSTableWriters.put(CObjectShardList.SHARD_INDEX_TABLE_NAME, Pair.create(this.buildSSTableWriterForShardIndexTable(sorted), (Map<CIndex, CQLSSTableWriter>) null));

	for (String defName : definitions.keySet()) {
		// Build the CQLSSTableWriter for the static table
		CQLSSTableWriter staticWriter = buildSSTableWriterForStaticTable(definitions.get(defName), sorted);

		// Build the CQLSSTableWriter for all the index tables
		List<CIndex> indexes = definitions.get(defName).getIndexesAsList();
		Map<CIndex, CQLSSTableWriter> indexWriters = Maps.newHashMap();
		for (CIndex index : indexes) {
			CQLSSTableWriter writer = buildSSTableWriterForWideTable(definitions.get(defName), index, sorted);
			indexWriters.put(index, writer);
		}
		this.SSTableWriters.put(defName, Pair.create(staticWriter, indexWriters));
	}
}
 

public void run(CommandLine cmdLine) throws Exception {
    Stopwatch stopwatch = Stopwatch.createStarted();

    if (cmdLine.hasOption("h")) {
        printUsage();
        return;
    }

    keyspace = cmdLine.getOptionValue("keyspace", "hawkular_metrics");
    dataDir = new File(cmdLine.getOptionValue("data-dir", "./data"));
    dataDir.mkdirs();
    tenants = Integer.parseInt(cmdLine.getOptionValue("tenants", "100"));
    metricsPerTenant = Integer.parseInt(cmdLine.getOptionValue("metrics-per-tenant", "100"));

    ValueServer valueServer = new ValueServer();
    valueServer.setMu(100);
    valueServer.setMode(ValueServer.UNIFORM_MODE);

    String endValue = cmdLine.getOptionValue("end");
    if (endValue == null) {
        endTime = System.currentTimeMillis();
    } else {
        endTime = getDuration("end", endValue, startEndRegexp);
    }

    String startValue = cmdLine.getOptionValue("start");
    if (startValue == null) {
        startTime = endTime - TimeUnit.MILLISECONDS.convert(1, TimeUnit.HOURS);
    } else {
        startTime = endTime - getDuration("start", startValue, startEndRegexp);
    }

    String intervalValue = cmdLine.getOptionValue("interval");
    if (intervalValue == null) {
        interval = TimeUnit.MILLISECONDS.convert(1, TimeUnit.MINUTES);
    } else {
        interval = getDuration("interval", intervalValue, intervalRegexp);
    }

    bufferSize = Integer.parseInt(cmdLine.getOptionValue("buffer-size", "128"));

    CQLSSTableWriter writer = createWriter();
    long totalDataPoints = 0;

    long currentTime = startTime;
    while (currentTime <= endTime) {
        for (int i = 0; i < tenants; ++i) {
            for (int j = 0; j < metricsPerTenant; ++j) {
                UUID timeUUID = TimeUUIDUtils.getTimeUUID(currentTime);
                writer.addRow("TENANT-" + i, GAUGE.getCode(), "GAUGE-" + j, 0L, timeUUID, valueServer.getNext());
                ++totalDataPoints;
            }
        }
        currentTime += interval;
    }

    writer.close();
    stopwatch.stop();

    System.out.println("\n\nStart time: " + startTime);
    System.out.println("End time: " + endTime);
    System.out.println("Total duration: " + (endTime - startTime) + " ms");
    System.out.println("Interval: " + interval);
    System.out.println("Tenants: " + tenants);
    System.out.println("Metrics per tenant: " + metricsPerTenant);
    System.out.println("Total data points: " + totalDataPoints);
    System.out.println("Execution time: " + stopwatch.elapsed(TimeUnit.MILLISECONDS) + " ms");
}
 

/**
 * Writes Rhombus objects into the appropriate static and index SSTableWriters for their object definition. Requires that initializeSSTableWriters
 * be called first and completeSSTableWrites be called when you're done inserting things. Object values with key "shardid" will be ignored and removed.
 * @param objects Map keyed by object name with a list of Rhombus objects to insert for that table
 * @throws CQLGenerationException
 * @throws IOException
 * @throws InvalidRequestException
 */
public void insertIntoSSTable(Map<String, List<Map<String, Object>>> objects) throws CQLGenerationException, IOException, InvalidRequestException {
	Map<String, CDefinition> definitions = this.keyspaceDefinition.getDefinitions();
	for (String tableName : objects.keySet()) {
		CDefinition definition = definitions.get(tableName);
		// Pull an existing CQLSSTableWriter for this object type if we already have one, if not make a new one
		if (!this.SSTableWriters.containsKey(tableName)) {
			throw new RuntimeException("Tried to write to uninitialized SSTableWriter for static table " + tableName);
		}
		CQLSSTableWriter staticWriter = this.SSTableWriters.get(tableName).left;
		Map<CIndex, CQLSSTableWriter> indexWriters = this.SSTableWriters.get(tableName).right;

		for (Map<String, Object> insert : objects.get(tableName)) {
			staticWriter.addRow(insert);
			for (CIndex index : indexWriters.keySet()) {
				if(definition.isAllowNullPrimaryKeyInserts()){
					//check if we have the necessary primary fields to insert on this index. If not just continue;
					if(!index.validateIndexKeys(index.getIndexKeyAndValues(insert))){
						continue;
					}
				}
				// Add the shard id to index writes
				long shardId = index.getShardingStrategy().getShardKey(insert.get("id"));
				insert.put("shardid", shardId);
				indexWriters.get(index).addRow(insert);
				// Pull the shardid back out to avoid the overhead of cloning the values and keep our abstraction from leaking
				insert.remove("shardid");

				// If this index uses shards, we need to record the write into the shard index table
				if((!(index.getShardingStrategy() instanceof ShardingStrategyNone))){
					String indexValuesString = CObjectCQLGenerator.makeIndexValuesString(index.getIndexValues(insert));
					Map<String, Object> shardIndexInsert = Maps.newHashMap();
					shardIndexInsert.put("tablename", CObjectCQLGenerator.makeTableName(definition, index));
					shardIndexInsert.put("indexvalues", indexValuesString);
					shardIndexInsert.put("shardid", shardId);
					shardIndexInsert.put("targetrowkey", shardId+":"+indexValuesString);
					this.SSTableWriters.get(CObjectShardList.SHARD_INDEX_TABLE_NAME).left.addRow(shardIndexInsert);
				}
			}
		}
	}
}