pkColumnsSet.set(pkSlotIndexesToBe[i] = i + posOffset);
targetColumns.set(i, indexIdColumn);
i++;
}
for (ColumnName colName : columnNodes) {
ColumnRef ref = resolver.resolveColumn(null, colName.getFamilyName(), colName.getColumnName());
PColumn column = ref.getColumn();
if (IndexUtil.getViewConstantValue(column, ptr)) {
if (overlapViewColumnsToBe.isEmpty()) {
overlapViewColumnsToBe = Sets.newHashSetWithExpectedSize(addViewColumnsToBe.size());
}
nColumnsToSet--;
overlapViewColumnsToBe.add(column);
addViewColumnsToBe.remove(column);
}
columnIndexesToBe[i] = ref.getColumnPosition();
targetColumns.set(i, column);
if (SchemaUtil.isPKColumn(column)) {
pkColumnsSet.set(pkSlotIndexesToBe[i] = ref.getPKSlotPosition());
}
i++;
}
for (PColumn column : addViewColumnsToBe) {
columnIndexesToBe[i] = column.getPosition();
targetColumns.set(i, column);
if (SchemaUtil.isPKColumn(column)) {
pkColumnsSet.set(pkSlotIndexesToBe[i] = SchemaUtil.getPKPosition(table, column));
}
i++;
}
for (i = posOffset; i < table.getPKColumns().size(); i++) {
PColumn pkCol = table.getPKColumns().get(i);
if (!pkColumnsSet.get(i)) {
if (!pkCol.isNullable()) {
throw new ConstraintViolationException(table.getName().getString() + "." + pkCol.getName().getString() + " may not be null");
}
}
}
}
boolean isAutoCommit = connection.getAutoCommit();
if (valueNodes == null) {
SelectStatement select = upsert.getSelect();
assert(select != null);
select = SubselectRewriter.flatten(select, connection);
ColumnResolver selectResolver = FromCompiler.getResolverForQuery(select, connection);
select = StatementNormalizer.normalize(select, selectResolver);
select = prependTenantAndViewConstants(table, select, tenantId, addViewColumnsToBe);
sameTable = select.getFrom().size() == 1
&& tableRefToBe.equals(selectResolver.getTables().get(0));
/* We can run the upsert in a coprocessor if:
* 1) from has only 1 table and the into table matches from table
* 2) the select query isn't doing aggregation (which requires a client-side final merge)
* 3) autoCommit is on
* 4) the table is not immutable with indexes, as the client is the one that figures out the additional
* puts for index tables.
* 5) no limit clause, as the limit clause requires client-side post processing
* 6) no sequences, as sequences imply that the order of upsert must match the order of
* selection.
* Otherwise, run the query to pull the data from the server
* and populate the MutationState (upto a limit).
*/
if (! (select.isAggregate() || select.isDistinct() || select.getLimit() != null || select.hasSequence()) ) {
// We can pipeline the upsert select instead of spooling everything to disk first,
// if we don't have any post processing that's required.
parallelIteratorFactoryToBe = new UpsertingParallelIteratorFactory(connection, tableRefToBe);
// If we're in the else, then it's not an aggregate, distinct, limted, or sequence using query,
// so we might be able to run it entirely on the server side.
runOnServer = sameTable && isAutoCommit && !(table.isImmutableRows() && !table.getIndexes().isEmpty());
}
// If we may be able to run on the server, add a hint that favors using the data table
// if all else is equal.
// TODO: it'd be nice if we could figure out in advance if the PK is potentially changing,
// as this would disallow running on the server. We currently use the row projector we
// get back to figure this out.
HintNode hint = upsert.getHint();
if (!upsert.getHint().hasHint(Hint.USE_INDEX_OVER_DATA_TABLE)) {
hint = HintNode.create(hint, Hint.USE_DATA_OVER_INDEX_TABLE);
}
select = SelectStatement.create(select, hint);
// Pass scan through if same table in upsert and select so that projection is computed correctly
// Use optimizer to choose the best plan
try {
QueryCompiler compiler = new QueryCompiler(statement, select, selectResolver, targetColumns, parallelIteratorFactoryToBe, new SequenceManager(statement));
queryPlanToBe = compiler.compile();
} catch (MetaDataEntityNotFoundException e) {
retryOnce = false; // don't retry if select clause has meta data entities that aren't found, as we already updated the cache
throw e;
}
nValuesToSet = queryPlanToBe.getProjector().getColumnCount();
// Cannot auto commit if doing aggregation or topN or salted
// Salted causes problems because the row may end up living on a different region
} else {
nValuesToSet = valueNodes.size() + addViewColumnsToBe.size() + (isTenantSpecific ? 1 : 0) + (isSharedViewIndex ? 1 : 0);
}
// Resize down to allow a subset of columns to be specifiable
if (columnNodes.isEmpty() && columnIndexesToBe.length >= nValuesToSet) {
nColumnsToSet = nValuesToSet;
columnIndexesToBe = Arrays.copyOf(columnIndexesToBe, nValuesToSet);
pkSlotIndexesToBe = Arrays.copyOf(pkSlotIndexesToBe, nValuesToSet);
}
if (nValuesToSet != nColumnsToSet) {
// We might have added columns, so refresh cache and try again if stale.
// Note that this check is not really sufficient, as a column could have
// been removed and the added back and we wouldn't detect that here.
if (retryOnce) {
retryOnce = false;
if (new MetaDataClient(connection).updateCache(schemaName, tableName).wasUpdated()) {
continue;
}
}
throw new SQLExceptionInfo.Builder(SQLExceptionCode.UPSERT_COLUMN_NUMBERS_MISMATCH)
.setMessage("Numbers of columns: " + nColumnsToSet + ". Number of values: " + nValuesToSet)
.build().buildException();
}
} catch (MetaDataEntityNotFoundException e) {
// Catch column/column family not found exception, as our meta data may
// be out of sync. Update the cache once and retry if we were out of sync.
// Otherwise throw, as we'll just get the same error next time.
if (retryOnce) {
retryOnce = false;
if (new MetaDataClient(connection).updateCache(schemaName, tableName).wasUpdated()) {
continue;
}
}
throw e;
}
break;
}
RowProjector projectorToBe = null;
// Optimize only after all checks have been performed
if (valueNodes == null) {
queryPlanToBe = new QueryOptimizer(services).optimize(queryPlanToBe, statement, targetColumns, parallelIteratorFactoryToBe);
projectorToBe = queryPlanToBe.getProjector();
runOnServer &= queryPlanToBe.getTableRef().equals(tableRefToBe);
}
final List<PColumn> allColumns = allColumnsToBe;
final RowProjector projector = projectorToBe;
final QueryPlan queryPlan = queryPlanToBe;
final TableRef tableRef = tableRefToBe;
final int[] columnIndexes = columnIndexesToBe;
final int[] pkSlotIndexes = pkSlotIndexesToBe;
final Set<PColumn> addViewColumns = addViewColumnsToBe;
final Set<PColumn> overlapViewColumns = overlapViewColumnsToBe;
final UpsertingParallelIteratorFactory parallelIteratorFactory = parallelIteratorFactoryToBe;
// TODO: break this up into multiple functions
////////////////////////////////////////////////////////////////////
// UPSERT SELECT
/////////////////////////////////////////////////////////////////////
if (valueNodes == null) {
// Before we re-order, check that for updatable view columns
// the projected expression either matches the column name or
// is a constant with the same required value.
throwIfNotUpdatable(tableRef, overlapViewColumnsToBe, targetColumns, projector, sameTable);
////////////////////////////////////////////////////////////////////
// UPSERT SELECT run server-side (maybe)
/////////////////////////////////////////////////////////////////////
if (runOnServer) {
// At most this array will grow bigger by the number of PK columns
int[] allColumnsIndexes = Arrays.copyOf(columnIndexes, columnIndexes.length + nValuesToSet);
int[] reverseColumnIndexes = new int[table.getColumns().size()];
List<Expression> projectedExpressions = Lists.newArrayListWithExpectedSize(reverseColumnIndexes.length);
Arrays.fill(reverseColumnIndexes, -1);
for (int i =0; i < nValuesToSet; i++) {
projectedExpressions.add(projector.getColumnProjector(i).getExpression());
reverseColumnIndexes[columnIndexes[i]] = i;
}
/*
* Order projected columns and projected expressions with PK columns
* leading order by slot position
*/
int offset = table.getBucketNum() == null ? 0 : 1;
for (int i = 0; i < table.getPKColumns().size() - offset; i++) {
PColumn column = table.getPKColumns().get(i + offset);
int pos = reverseColumnIndexes[column.getPosition()];
if (pos == -1) {
// Last PK column may be fixed width and nullable
// We don't want to insert a null expression b/c
// it's not valid to set a fixed width type to null.
if (column.getDataType().isFixedWidth()) {
continue;
}
// Add literal null for missing PK columns
pos = projectedExpressions.size();
Expression literalNull = LiteralExpression.newConstant(null, column.getDataType(), true);
projectedExpressions.add(literalNull);
allColumnsIndexes[pos] = column.getPosition();
}
// Swap select expression at pos with i
Collections.swap(projectedExpressions, i, pos);
// Swap column indexes and reverse column indexes too
int tempPos = allColumnsIndexes[i];
allColumnsIndexes[i] = allColumnsIndexes[pos];
allColumnsIndexes[pos] = tempPos;
reverseColumnIndexes[tempPos] = reverseColumnIndexes[i];
reverseColumnIndexes[i] = i;
}
// If any pk slots are changing, be conservative and don't run this server side.
// If the row ends up living in a different region, we'll get an error otherwise.
for (int i = 0; i < table.getPKColumns().size(); i++) {
PColumn column = table.getPKColumns().get(i);
Expression source = projectedExpressions.get(i);
if (source == null || !source.equals(new ColumnRef(tableRef, column.getPosition()).newColumnExpression())) {
// TODO: we could check the region boundaries to see if the pk will still be in it.
runOnServer = false; // bail on running server side, since PK may be changing
break;
}
}