package org.h2.command.dml;
import java.util.ArrayList;
import java.util.HashSet;
import org.h2.api.ErrorCode;
import org.h2.engine.Database;
import org.h2.engine.Mode;
import org.h2.engine.Session;
import org.h2.expression.Expression;
import org.h2.expression.ExpressionColumn;
import org.h2.expression.ExpressionVisitor;
import org.h2.expression.Parameter;
import org.h2.expression.ValueExpression;
import org.h2.message.DbException;
import org.h2.result.LazyResult;
import org.h2.result.LocalResult;
import org.h2.result.ResultInterface;
import org.h2.result.ResultTarget;
import org.h2.table.Column;
import org.h2.table.ColumnResolver;
import org.h2.table.Table;
import org.h2.table.TableFilter;
import org.h2.util.ColumnNamer;
import org.h2.value.Value;
import org.h2.value.ValueInt;
import org.h2.value.ValueNull;
public class SelectUnion extends Query {
public enum UnionType {
UNION,
UNION_ALL,
EXCEPT,
INTERSECT
}
private final UnionType unionType;
final Query left;
final Query right;
private boolean isPrepared, checkInit;
private boolean isForUpdate;
public SelectUnion(Session session, UnionType unionType, Query query, Query right) {
super(session);
this.unionType = unionType;
this.left = query;
this.right = right;
}
@Override
public boolean isUnion() {
return true;
}
@Override
public void prepareJoinBatch() {
left.prepareJoinBatch();
right.prepareJoinBatch();
}
public UnionType getUnionType() {
return unionType;
}
public Query getLeft() {
return left;
}
public Query getRight() {
return right;
}
@Override
public void setDistinctIfPossible() {
setDistinct();
}
private Value[] convert(Value[] values, int columnCount) {
Value[] newValues;
if (columnCount == values.length) {
newValues = values;
} else {
newValues = new Value[columnCount];
}
Mode mode = session.getDatabase().getMode();
for (int i = 0; i < columnCount; i++) {
Expression e = expressions.get(i);
newValues[i] = values[i].convertTo(e.getType(), mode, null);
}
return newValues;
}
@Override
public ResultInterface queryMeta() {
int columnCount = left.getColumnCount();
LocalResult result = session.getDatabase().getResultFactory().create(session, expressionArray, columnCount);
result.done();
return result;
}
public LocalResult getEmptyResult() {
int columnCount = left.getColumnCount();
return session.getDatabase().getResultFactory().create(session, expressionArray, columnCount);
}
@Override
protected ResultInterface queryWithoutCache(int maxRows, ResultTarget target) {
if (maxRows != 0) {
int l;
if (limitExpr == null) {
l = -1;
} else {
Value v = limitExpr.getValue(session);
l = v == ValueNull.INSTANCE ? -1 : v.getInt();
}
if (l < 0) {
l = maxRows;
} else {
l = Math.min(l, maxRows);
}
limitExpr = ValueExpression.get(ValueInt.get(l));
}
Database db = session.getDatabase();
if (db.getSettings().optimizeInsertFromSelect) {
if (unionType == UnionType.UNION_ALL && target != null) {
if (sort == null && !distinct && maxRows == 0 &&
offsetExpr == null && limitExpr == null) {
left.query(0, target);
right.query(0, target);
return null;
}
}
}
int columnCount = left.getColumnCount();
if (session.isLazyQueryExecution() && unionType == UnionType.UNION_ALL && !distinct &&
sort == null && !randomAccessResult && !isForUpdate &&
offsetExpr == null && !fetchPercent && !withTies && isReadOnly()) {
int limit = -1;
if (limitExpr != null) {
Value v = limitExpr.getValue(session);
if (v != ValueNull.INSTANCE) {
limit = v.getInt();
}
}
if (limit != 0) {
LazyResultUnion lazyResult = new LazyResultUnion(expressionArray, columnCount);
if (limit > 0) {
lazyResult.setLimit(limit);
}
return lazyResult;
}
}
LocalResult result = db.getResultFactory().create(session, expressionArray, columnCount);
if (sort != null) {
result.setSortOrder(sort);
}
if (distinct) {
left.setDistinctIfPossible();
right.setDistinctIfPossible();
result.setDistinct();
}
switch (unionType) {
case UNION:
case EXCEPT:
left.setDistinctIfPossible();
right.setDistinctIfPossible();
result.setDistinct();
break;
case UNION_ALL:
break;
case INTERSECT:
left.setDistinctIfPossible();
right.setDistinctIfPossible();
break;
default:
DbException.throwInternalError("type=" + unionType);
}
ResultInterface l = left.query(0);
ResultInterface r = right.query(0);
l.reset();
r.reset();
switch (unionType) {
case UNION_ALL:
case UNION: {
while (l.next()) {
result.addRow(convert(l.currentRow(), columnCount));
}
while (r.next()) {
result.addRow(convert(r.currentRow(), columnCount));
}
break;
}
case EXCEPT: {
while (l.next()) {
result.addRow(convert(l.currentRow(), columnCount));
}
while (r.next()) {
result.removeDistinct(convert(r.currentRow(), columnCount));
}
break;
}
case INTERSECT: {
LocalResult temp = db.getResultFactory().create(session, expressionArray, columnCount);
temp.setDistinct();
while (l.next()) {
temp.addRow(convert(l.currentRow(), columnCount));
}
while (r.next()) {
Value[] values = convert(r.currentRow(), columnCount);
if (temp.containsDistinct(values)) {
result.addRow(values);
}
}
temp.close();
break;
}
default:
DbException.throwInternalError("type=" + unionType);
}
if (offsetExpr != null) {
result.setOffset(offsetExpr.getValue(session).getInt());
}
if (limitExpr != null) {
Value v = limitExpr.getValue(session);
if (v != ValueNull.INSTANCE) {
result.setLimit(v.getInt());
result.setFetchPercent(fetchPercent);
if (withTies) {
result.setWithTies(sort);
}
}
}
l.close();
r.close();
result.done();
if (target != null) {
while (result.next()) {
target.addRow(result.currentRow());
}
result.close();
return null;
}
return result;
}
@Override
public void init() {
if (checkInit) {
DbException.throwInternalError();
}
checkInit = true;
left.init();
right.init();
int len = left.getColumnCount();
if (len != right.getColumnCount()) {
throw DbException.get(ErrorCode.COLUMN_COUNT_DOES_NOT_MATCH);
}
ArrayList<Expression> le = left.getExpressions();
expressions = new ArrayList<>(len);
for (int i = 0; i < len; i++) {
Expression l = le.get(i);
expressions.add(l);
}
if (withTies && !hasOrder()) {
throw DbException.get(ErrorCode.WITH_TIES_WITHOUT_ORDER_BY);
}
}
@Override
public void prepare() {
if (isPrepared) {
return;
}
if (!checkInit) {
DbException.throwInternalError("not initialized");
}
isPrepared = true;
left.prepare();
right.prepare();
int len = left.getColumnCount();
expressions = new ArrayList<>(len);
ArrayList<Expression> le = left.getExpressions();
ArrayList<Expression> re = right.getExpressions();
ColumnNamer columnNamer= new ColumnNamer(session);
for (int i = 0; i < len; i++) {
Expression l = le.get(i);
Expression r = re.get(i);
String columnName = columnNamer.getColumnName(l, i, l.getAlias());
Column col = new Column(columnName, Value.getHigherType(l.getType(), r.getType()));
Expression e = new ExpressionColumn(session.getDatabase(), col);
expressions.add(e);
}
if (orderList != null) {
initOrder(session, expressions, null, orderList, getColumnCount(), true, null);
sort = prepareOrder(orderList, expressions.size());
orderList = null;
}
expressionArray = expressions.toArray(new Expression[0]);
}
@Override
public double getCost() {
return left.getCost() + right.getCost();
}
@Override
public HashSet<Table> getTables() {
HashSet<Table> set = left.getTables();
set.addAll(right.getTables());
return set;
}
@Override
public void setForUpdate(boolean forUpdate) {
left.setForUpdate(forUpdate);
right.setForUpdate(forUpdate);
isForUpdate = forUpdate;
}
@Override
public int getColumnCount() {
return left.getColumnCount();
}
@Override
public void mapColumns(ColumnResolver resolver, int level) {
left.mapColumns(resolver, level);
right.mapColumns(resolver, level);
}
@Override
public void setEvaluatable(TableFilter tableFilter, boolean b) {
left.setEvaluatable(tableFilter, b);
right.setEvaluatable(tableFilter, b);
}
@Override
public void addGlobalCondition(Parameter param, int columnId,
int comparisonType) {
addParameter(param);
switch (unionType) {
case UNION_ALL:
case UNION:
case INTERSECT: {
left.addGlobalCondition(param, columnId, comparisonType);
right.addGlobalCondition(param, columnId, comparisonType);
break;
}
case EXCEPT: {
left.addGlobalCondition(param, columnId, comparisonType);
break;
}
default:
DbException.throwInternalError("type=" + unionType);
}
}
@Override
public String getPlanSQL(boolean alwaysQuote) {
StringBuilder buff = new StringBuilder();
buff.append('(').append(left.getPlanSQL(alwaysQuote)).append(')');
switch (unionType) {
case UNION_ALL:
buff.append("\nUNION ALL\n");
break;
case UNION:
buff.append("\nUNION\n");
break;
case INTERSECT:
buff.append("\nINTERSECT\n");
break;
case EXCEPT:
buff.append("\nEXCEPT\n");
break;
default:
DbException.throwInternalError("type=" + unionType);
}
buff.append('(').append(right.getPlanSQL(alwaysQuote)).append(')');
Expression[] exprList = expressions.toArray(new Expression[0]);
if (sort != null) {
buff.append("\nORDER BY ").append(sort.getSQL(exprList, exprList.length, alwaysQuote));
}
appendLimitToSQL(buff, alwaysQuote);
if (sampleSizeExpr != null) {
buff.append("\nSAMPLE_SIZE ");
sampleSizeExpr.getUnenclosedSQL(buff, alwaysQuote);
}
if (isForUpdate) {
buff.append("\nFOR UPDATE");
}
return buff.toString();
}
@Override
public boolean isEverything(ExpressionVisitor visitor) {
return left.isEverything(visitor) && right.isEverything(visitor);
}
@Override
public boolean isReadOnly() {
return left.isReadOnly() && right.isReadOnly();
}
@Override
public void updateAggregate(Session s, int stage) {
left.updateAggregate(s, stage);
right.updateAggregate(s, stage);
}
@Override
public void fireBeforeSelectTriggers() {
left.fireBeforeSelectTriggers();
right.fireBeforeSelectTriggers();
}
@Override
public boolean allowGlobalConditions() {
return left.allowGlobalConditions() && right.allowGlobalConditions();
}
private final class LazyResultUnion extends LazyResult {
int columnCount;
ResultInterface l;
ResultInterface r;
boolean leftDone;
boolean rightDone;
LazyResultUnion(Expression[] expressions, int columnCount) {
super(expressions);
this.columnCount = columnCount;
}
@Override
public int getVisibleColumnCount() {
return columnCount;
}
@Override
protected Value[] fetchNextRow() {
if (rightDone) {
return null;
}
if (!leftDone) {
if (l == null) {
l = left.query(0);
l.reset();
}
if (l.next()) {
return l.currentRow();
}
leftDone = true;
}
if (r == null) {
r = right.query(0);
r.reset();
}
if (r.next()) {
return r.currentRow();
}
rightDone = true;
return null;
}
@Override
public void close() {
super.close();
if (l != null) {
l.close();
}
if (r != null) {
r.close();
}
}
@Override
public void reset() {
super.reset();
if (l != null) {
l.reset();
}
if (r != null) {
r.reset();
}
leftDone = false;
rightDone = false;
}
}
}