[enhancement](Nereids) Speedup PartitionPrunner (#31970)

This pr imporve the high QPS query by speed up PartitionPrunner

1. remove useless Date parse/format, use LocalDate instead
2. fast evaluate path for single value partition
3. change Collection.stream() to ImmutableXxx.builderWithExpectedSize(n) to skip useless method call and collection resize
4. change lots of if-else to switch
5. don't parse to string to compare dateLiteral, use int field compare instead
This commit is contained in:
924060929
2024-03-08 13:38:57 +08:00
committed by yiguolei
parent e3611f6a1d
commit e8b4bf5be9
8 changed files with 401 additions and 251 deletions

View File

@ -193,6 +193,14 @@ public class DateLiteral extends LiteralExpr {
private static final Pattern HAS_OFFSET_PART = Pattern.compile("[\\+\\-]\\d{2}:\\d{2}");
@Override
public boolean equals(Object o) {
if (o instanceof DateLiteral) {
return compareLiteral((LiteralExpr) o) == 0;
}
return super.equals(o);
}
// Date Literal persist type in meta
private enum DateLiteralType {
DATETIME(0),
@ -626,6 +634,31 @@ public class DateLiteral extends LiteralExpr {
if (expr == MaxLiteral.MAX_VALUE) {
return -1;
}
if (expr instanceof DateLiteral) {
DateLiteral other = (DateLiteral) expr;
long yearMonthDay = year * 10000 + month * 100 + day;
long otherYearMonthDay = other.year * 10000 + other.month * 100 + other.day;
long diffDay = yearMonthDay - otherYearMonthDay;
if (diffDay != 0) {
return diffDay < 0 ? -1 : 1;
}
int typeAsInt = isDateType() ? 0 : 1;
int thatTypeAsInt = other.isDateType() ? 0 : 1;
int typeDiff = typeAsInt - thatTypeAsInt;
if (typeDiff != 0) {
return typeDiff;
}
long hourMinuteSecond = hour * 10000 + minute * 100 + second;
long otherHourMinuteSecond = other.hour * 10000 + other.minute * 100 + other.second;
long diffSecond = hourMinuteSecond - otherHourMinuteSecond;
if (diffSecond != 0) {
return diffSecond < 0 ? -1 : 1;
}
long diff = getMicroPartWithinScale() - other.getMicroPartWithinScale();
return diff < 0 ? -1 : (diff == 0 ? 0 : 1);
}
// date time will not overflow when doing addition and subtraction
return getStringValue().compareTo(expr.getStringValue());
}
@ -1731,6 +1764,15 @@ public class DateLiteral extends LiteralExpr {
throw new InvalidFormatException("'" + value + "' is invalid");
}
private long getMicroPartWithinScale() {
if (type.isDatetimeV2()) {
int scale = ((ScalarType) type).getScalarScale();
return (long) (microsecond / SCALE_FACTORS[scale]);
} else {
return 0;
}
}
public void setMinValue() {
year = 0;
month = 1;

View File

@ -53,6 +53,7 @@ import org.apache.doris.nereids.util.Utils;
import com.google.common.collect.BoundType;
import com.google.common.collect.ImmutableList;
import com.google.common.collect.ImmutableMap;
import com.google.common.collect.ImmutableMap.Builder;
import com.google.common.collect.ImmutableSet;
import com.google.common.collect.Lists;
import com.google.common.collect.Maps;
@ -62,11 +63,11 @@ import java.util.ArrayList;
import java.util.LinkedHashMap;
import java.util.List;
import java.util.Map;
import java.util.Map.Entry;
import java.util.Objects;
import java.util.Optional;
import java.util.Set;
import java.util.function.BiFunction;
import java.util.stream.IntStream;
/**
* OneRangePartitionEvaluator.
@ -102,31 +103,33 @@ public class OneRangePartitionEvaluator
PartitionRangeExpander expander = new PartitionRangeExpander();
this.partitionSlotTypes = expander.computePartitionSlotTypes(lowers, uppers);
this.slotToType = IntStream.range(0, partitionSlots.size())
.mapToObj(index -> Pair.of(partitionSlots.get(index), partitionSlotTypes.get(index)))
.collect(ImmutableMap.toImmutableMap(Pair::key, Pair::value));
this.slotToType = Maps.newHashMapWithExpectedSize(partitionSlots.size() * 2);
for (int i = 0; i < partitionSlots.size(); i++) {
slotToType.put(partitionSlots.get(i), partitionSlotTypes.get(i));
}
this.partitionSlotContainsNull = IntStream.range(0, partitionSlots.size())
.mapToObj(index -> {
Slot slot = partitionSlots.get(index);
if (!slot.nullable()) {
return Pair.of(slot, false);
}
PartitionSlotType partitionSlotType = partitionSlotTypes.get(index);
boolean maybeNull = false;
switch (partitionSlotType) {
case CONST:
case RANGE:
maybeNull = range.lowerEndpoint().getKeys().get(index).isMinValue();
break;
case OTHER:
maybeNull = true;
break;
default:
throw new AnalysisException("Unknown partition slot type: " + partitionSlotType);
}
return Pair.of(slot, maybeNull);
}).collect(ImmutableMap.toImmutableMap(Pair::key, Pair::value));
this.partitionSlotContainsNull = Maps.newHashMapWithExpectedSize(partitionSlots.size() * 2);
for (int i = 0; i < partitionSlots.size(); i++) {
Slot slot = partitionSlots.get(i);
if (!slot.nullable()) {
partitionSlotContainsNull.put(slot, false);
continue;
}
PartitionSlotType partitionSlotType = partitionSlotTypes.get(i);
boolean maybeNull = false;
switch (partitionSlotType) {
case CONST:
case RANGE:
maybeNull = range.lowerEndpoint().getKeys().get(i).isMinValue();
break;
case OTHER:
maybeNull = true;
break;
default:
throw new AnalysisException("Unknown partition slot type: " + partitionSlotType);
}
partitionSlotContainsNull.put(slot, maybeNull);
}
int expandThreshold = cascadesContext.getAndCacheSessionVariable(
"partitionPruningExpandThreshold",
@ -147,62 +150,14 @@ public class OneRangePartitionEvaluator
@Override
public List<Map<Slot, PartitionSlotInput>> getOnePartitionInputs() {
List<Map<Slot, PartitionSlotInput>> onePartitionInputs = Lists.newArrayList();
for (List<Expression> input : inputs) {
boolean previousIsLowerBoundLiteral = true;
boolean previousIsUpperBoundLiteral = true;
List<Pair<Slot, PartitionSlotInput>> slotToInputs = Lists.newArrayList();
for (int i = 0; i < partitionSlots.size(); ++i) {
Slot partitionSlot = partitionSlots.get(i);
// partitionSlot will be replaced to this expression
Expression expression = input.get(i);
ColumnRange slotRange = null;
PartitionSlotType partitionSlotType = partitionSlotTypes.get(i);
if (expression instanceof Literal) {
// const or expanded range
slotRange = ColumnRange.singleton((Literal) expression);
if (!expression.equals(lowers.get(i))) {
previousIsLowerBoundLiteral = false;
}
if (!expression.equals(uppers.get(i))) {
previousIsUpperBoundLiteral = false;
}
} else {
// un expanded range
switch (partitionSlotType) {
case RANGE:
boolean isLastPartitionColumn = i + 1 == partitionSlots.size();
BoundType rightBoundType = isLastPartitionColumn
? BoundType.OPEN : BoundType.CLOSED;
slotRange = ColumnRange.range(
lowers.get(i), BoundType.CLOSED, uppers.get(i), rightBoundType);
break;
case OTHER:
if (previousIsLowerBoundLiteral) {
slotRange = ColumnRange.atLeast(lowers.get(i));
} else if (previousIsUpperBoundLiteral) {
slotRange = ColumnRange.lessThen(uppers.get(i));
} else {
// unknown range
slotRange = ColumnRange.all();
}
break;
default:
throw new AnalysisException("Unknown partition slot type: " + partitionSlotType);
}
previousIsLowerBoundLiteral = false;
previousIsUpperBoundLiteral = false;
}
ImmutableMap<Slot, ColumnRange> slotToRange = ImmutableMap.of(partitionSlot, slotRange);
slotToInputs.add(Pair.of(partitionSlot, new PartitionSlotInput(expression, slotToRange)));
}
Map<Slot, PartitionSlotInput> slotPartitionSlotInputMap = fillSlotRangesToInputs(
slotToInputs.stream()
.collect(ImmutableMap.toImmutableMap(Pair::key, Pair::value)));
onePartitionInputs.add(slotPartitionSlotInputMap);
if (partitionSlots.size() == 1 && inputs.size() == 1 && inputs.get(0).size() == 1
&& inputs.get(0).get(0) instanceof Literal) {
// fast path
return computeSinglePartitionValueInputs();
} else {
// slow path
return commonComputeOnePartitionInputs();
}
return onePartitionInputs;
}
@Override
@ -597,13 +552,14 @@ public class OneRangePartitionEvaluator
}
private List<Literal> toNereidsLiterals(PartitionKey partitionKey) {
return IntStream.range(0, partitionKey.getKeys().size())
.mapToObj(index -> {
LiteralExpr literalExpr = partitionKey.getKeys().get(index);
PrimitiveType primitiveType = partitionKey.getTypes().get(index);
Type type = Type.fromPrimitiveType(primitiveType);
return Literal.fromLegacyLiteral(literalExpr, type);
}).collect(ImmutableList.toImmutableList());
List<Literal> literals = Lists.newArrayListWithCapacity(partitionKey.getKeys().size());
for (int i = 0; i < partitionKey.getKeys().size(); i++) {
LiteralExpr literalExpr = partitionKey.getKeys().get(i);
PrimitiveType primitiveType = partitionKey.getTypes().get(i);
Type type = Type.fromPrimitiveType(primitiveType);
literals.add(Literal.fromLegacyLiteral(literalExpr, type));
}
return literals;
}
@Override
@ -655,15 +611,20 @@ public class OneRangePartitionEvaluator
private Map<Slot, PartitionSlotInput> fillSlotRangesToInputs(
Map<Slot, PartitionSlotInput> inputs) {
Map<Slot, ColumnRange> allColumnRanges = inputs.entrySet()
.stream()
.map(entry -> Pair.of(entry.getKey(), entry.getValue().columnRanges.get(entry.getKey())))
.collect(ImmutableMap.toImmutableMap(Pair::key, Pair::value));
Builder<Slot, ColumnRange> allColumnRangesBuilder =
ImmutableMap.builderWithExpectedSize(inputs.size() * 2);
for (Entry<Slot, PartitionSlotInput> entry : inputs.entrySet()) {
allColumnRangesBuilder.put(entry.getKey(), entry.getValue().columnRanges.get(entry.getKey()));
}
return inputs.keySet()
.stream()
.map(slot -> Pair.of(slot, new PartitionSlotInput(inputs.get(slot).result, allColumnRanges)))
.collect(ImmutableMap.toImmutableMap(Pair::key, Pair::value));
Map<Slot, ColumnRange> allColumnRanges = allColumnRangesBuilder.build();
Builder<Slot, PartitionSlotInput> partitionSlotInputs =
ImmutableMap.builderWithExpectedSize(inputs.size() * 2);
for (Slot slot : inputs.keySet()) {
partitionSlotInputs.put(slot, new PartitionSlotInput(inputs.get(slot).result, allColumnRanges));
}
return partitionSlotInputs.build();
}
/** EvaluateRangeInput */
@ -728,4 +689,71 @@ public class OneRangePartitionEvaluator
public boolean isDefaultPartition() {
return partitionItem.isDefaultPartition();
}
private List<Map<Slot, PartitionSlotInput>> computeSinglePartitionValueInputs() {
Slot partitionSlot = partitionSlots.get(0);
Literal literal = (Literal) inputs.get(0).get(0);
ColumnRange slotRange = ColumnRange.singleton(literal);
ImmutableMap<Slot, ColumnRange> slotToRange = ImmutableMap.of(partitionSlot, slotRange);
Map<Slot, PartitionSlotInput> slotToInputs =
ImmutableMap.of(partitionSlot, new PartitionSlotInput(literal, slotToRange));
return ImmutableList.of(slotToInputs);
}
private List<Map<Slot, PartitionSlotInput>> commonComputeOnePartitionInputs() {
List<Map<Slot, PartitionSlotInput>> onePartitionInputs = Lists.newArrayListWithCapacity(inputs.size());
for (List<Expression> input : inputs) {
boolean previousIsLowerBoundLiteral = true;
boolean previousIsUpperBoundLiteral = true;
Builder<Slot, PartitionSlotInput> slotToInputs = ImmutableMap.builderWithExpectedSize(16);
for (int i = 0; i < partitionSlots.size(); ++i) {
Slot partitionSlot = partitionSlots.get(i);
// partitionSlot will be replaced to this expression
Expression expression = input.get(i);
ColumnRange slotRange = null;
PartitionSlotType partitionSlotType = partitionSlotTypes.get(i);
if (expression instanceof Literal) {
// const or expanded range
slotRange = ColumnRange.singleton((Literal) expression);
if (!expression.equals(lowers.get(i))) {
previousIsLowerBoundLiteral = false;
}
if (!expression.equals(uppers.get(i))) {
previousIsUpperBoundLiteral = false;
}
} else {
// un expanded range
switch (partitionSlotType) {
case RANGE:
boolean isLastPartitionColumn = i + 1 == partitionSlots.size();
BoundType rightBoundType = isLastPartitionColumn
? BoundType.OPEN : BoundType.CLOSED;
slotRange = ColumnRange.range(
lowers.get(i), BoundType.CLOSED, uppers.get(i), rightBoundType);
break;
case OTHER:
if (previousIsLowerBoundLiteral) {
slotRange = ColumnRange.atLeast(lowers.get(i));
} else if (previousIsUpperBoundLiteral) {
slotRange = ColumnRange.lessThen(uppers.get(i));
} else {
// unknown range
slotRange = ColumnRange.all();
}
break;
default:
throw new AnalysisException("Unknown partition slot type: " + partitionSlotType);
}
previousIsLowerBoundLiteral = false;
previousIsUpperBoundLiteral = false;
}
ImmutableMap<Slot, ColumnRange> slotToRange = ImmutableMap.of(partitionSlot, slotRange);
slotToInputs.put(partitionSlot, new PartitionSlotInput(expression, slotToRange));
}
Map<Slot, PartitionSlotInput> slotPartitionSlotInputMap = fillSlotRangesToInputs(slotToInputs.build());
onePartitionInputs.add(slotPartitionSlotInputMap);
}
return onePartitionInputs;
}
}

View File

@ -34,10 +34,13 @@ import org.apache.doris.nereids.trees.expressions.visitor.DefaultExpressionRewri
import org.apache.doris.nereids.types.DateTimeType;
import com.google.common.collect.ImmutableList;
import com.google.common.collect.ImmutableList.Builder;
import com.google.common.collect.ImmutableSet;
import com.google.common.collect.Lists;
import java.util.List;
import java.util.Map;
import java.util.Map.Entry;
import java.util.Objects;
/**
@ -91,11 +94,15 @@ public class PartitionPruner extends DefaultExpressionRewriter<Void> {
}
}
/** prune */
public List<Long> prune() {
return partitions.stream()
.filter(partitionEvaluator -> !canPrune(partitionEvaluator))
.map(OnePartitionEvaluator::getPartitionId)
.collect(ImmutableList.toImmutableList());
Builder<Long> scanPartitionIds = ImmutableList.builder();
for (OnePartitionEvaluator partition : partitions) {
if (!canPrune(partition)) {
scanPartitionIds.add(partition.getPartitionId());
}
}
return scanPartitionIds.build();
}
/**
@ -107,11 +114,12 @@ public class PartitionPruner extends DefaultExpressionRewriter<Void> {
partitionPredicate = TryEliminateUninterestedPredicates.rewrite(
partitionPredicate, ImmutableSet.copyOf(partitionSlots), cascadesContext);
partitionPredicate = PredicateRewriteForPartitionPrune.rewrite(partitionPredicate, cascadesContext);
List<OnePartitionEvaluator> evaluators = idToPartitions.entrySet()
.stream()
.map(kv -> toPartitionEvaluator(kv.getKey(), kv.getValue(), partitionSlots, cascadesContext,
partitionTableType))
.collect(ImmutableList.toImmutableList());
List<OnePartitionEvaluator> evaluators = Lists.newArrayListWithCapacity(idToPartitions.size());
for (Entry<Long, PartitionItem> kv : idToPartitions.entrySet()) {
evaluators.add(toPartitionEvaluator(
kv.getKey(), kv.getValue(), partitionSlots, cascadesContext, partitionTableType));
}
partitionPredicate = OrToIn.INSTANCE.rewrite(partitionPredicate, null);
PartitionPruner partitionPruner = new PartitionPruner(evaluators, partitionPredicate);

View File

@ -34,13 +34,11 @@ import org.apache.doris.nereids.types.DataType;
import com.google.common.collect.ImmutableList;
import com.google.common.collect.Iterators;
import com.google.common.collect.Lists;
import org.apache.commons.lang3.time.DateFormatUtils;
import org.apache.commons.lang3.time.DateUtils;
import org.apache.commons.collections.iterators.SingletonIterator;
import java.math.BigInteger;
import java.text.ParseException;
import java.time.temporal.ChronoUnit;
import java.util.Date;
import java.time.LocalDate;
import java.time.ZoneOffset;
import java.util.Iterator;
import java.util.List;
import java.util.NoSuchElementException;
@ -56,6 +54,8 @@ import java.util.function.Function;
* after expand range, we can replace partition slot to the literal in expression tree and evaluate it.
*/
public class PartitionRangeExpander {
public static final long ONE_DAY_MILLIS_SECOND = 1000L * 60 * 60 * 24;
/** PartitionSlotType */
public enum PartitionSlotType {
// e.g. the first partition column is const '1' in partition [('1', '2', '5'), ('1', '3', '5')),
@ -83,7 +83,7 @@ public class PartitionRangeExpander {
for (int i = 0; i < partitionSlotTypes.size(); i++) {
Slot slot = partitionSlots.get(i);
PartitionSlotType partitionSlotType = partitionSlotTypes.get(i);
List<Expression> expandedList = Lists.newArrayList();
List<Expression> expandedList = Lists.newArrayListWithCapacity(2);
Literal lower = lowers.get(i);
switch (partitionSlotType) {
case CONST:
@ -98,9 +98,13 @@ public class PartitionRangeExpander {
try {
boolean isLastColumn = i + 1 == partitionSlots.size();
if (canExpandRange(slot, lower, upper, expandedCount, expandThreshold)) {
expandedList.addAll(ImmutableList.copyOf(
enumerableIterator(slot, lower, upper, isLastColumn))
);
Iterator<? extends Expression> iterator = enumerableIterator(
slot, lower, upper, isLastColumn);
if (iterator instanceof SingletonIterator) {
expandedList.add(iterator.next());
} else {
expandedList.addAll(ImmutableList.copyOf(iterator));
}
} else {
expandedList.add(slot);
}
@ -122,7 +126,7 @@ public class PartitionRangeExpander {
return expandedLists;
}
private final boolean canExpandRange(Slot slot, Literal lower, Literal upper,
private boolean canExpandRange(Slot slot, Literal lower, Literal upper,
long expandedCount, int expandThreshold) {
DataType type = slot.getDataType();
if (!type.isIntegerLikeType() && !type.isDateType() && !type.isDateV2Type()) {
@ -163,24 +167,53 @@ public class PartitionRangeExpander {
return types;
}
private final long enumerableCount(DataType dataType, Literal startInclusive, Literal endExclusive) throws
ParseException {
private long enumerableCount(DataType dataType, Literal startInclusive, Literal endExclusive) {
if (dataType.isIntegerLikeType()) {
BigInteger start = new BigInteger(startInclusive.getStringValue());
BigInteger end = new BigInteger(endExclusive.getStringValue());
return end.subtract(start).longValue();
} else if (dataType.isDateType() || dataType.isDateV2Type()) {
Date start = DateUtils.parseDate(startInclusive.toString(), DateLiteral.JAVA_DATE_FORMAT);
Date end = DateUtils.parseDate(endExclusive.toString(), DateLiteral.JAVA_DATE_FORMAT);
return ChronoUnit.DAYS.between(start.toInstant(), end.toInstant());
} else if (dataType.isDateType()) {
DateLiteral startInclusiveDate = (DateLiteral) startInclusive;
DateLiteral endExclusiveDate = (DateLiteral) endExclusive;
LocalDate startDate = LocalDate.of(
(int) startInclusiveDate.getYear(),
(int) startInclusiveDate.getMonth(),
(int) startInclusiveDate.getDay()
);
LocalDate endDate = LocalDate.of(
(int) endExclusiveDate.getYear(),
(int) endExclusiveDate.getMonth(),
(int) endExclusiveDate.getDay()
);
long diffMillisSecond = endDate.atStartOfDay(ZoneOffset.UTC).toInstant().toEpochMilli()
- startDate.atStartOfDay(ZoneOffset.UTC).toInstant().toEpochMilli();
return (diffMillisSecond + ONE_DAY_MILLIS_SECOND + ONE_DAY_MILLIS_SECOND - 1) / ONE_DAY_MILLIS_SECOND;
} else if (dataType.isDateV2Type()) {
DateV2Literal startInclusiveDate = (DateV2Literal) startInclusive;
DateV2Literal endExclusiveDate = (DateV2Literal) endExclusive;
LocalDate startDate = LocalDate.of(
(int) startInclusiveDate.getYear(),
(int) startInclusiveDate.getMonth(),
(int) startInclusiveDate.getDay()
);
LocalDate endDate = LocalDate.of(
(int) endExclusiveDate.getYear(),
(int) endExclusiveDate.getMonth(),
(int) endExclusiveDate.getDay()
);
long diffMillisSecond = endDate.atStartOfDay(ZoneOffset.UTC).toInstant().toEpochMilli()
- startDate.atStartOfDay(ZoneOffset.UTC).toInstant().toEpochMilli();
return (diffMillisSecond + ONE_DAY_MILLIS_SECOND + ONE_DAY_MILLIS_SECOND - 1) / ONE_DAY_MILLIS_SECOND;
}
// not enumerable
return -1;
}
private final Iterator<? extends Expression> enumerableIterator(
Slot slot, Literal startInclusive, Literal endLiteral, boolean endExclusive) throws ParseException {
private Iterator<? extends Expression> enumerableIterator(
Slot slot, Literal startInclusive, Literal endLiteral, boolean endExclusive) {
DataType dataType = slot.getDataType();
if (dataType.isIntegerLikeType()) {
BigInteger start = new BigInteger(startInclusive.getStringValue());
@ -202,15 +235,48 @@ public class PartitionRangeExpander {
start, end, endExclusive, LargeIntLiteral::new);
}
} else if (dataType.isDateType()) {
Date startDate = DateUtils.parseDate(startInclusive.toString(), DateLiteral.JAVA_DATE_FORMAT);
Date endDate = DateUtils.parseDate(endLiteral.toString(), DateLiteral.JAVA_DATE_FORMAT);
DateLiteral startInclusiveDate = (DateLiteral) startInclusive;
DateLiteral endLiteralDate = (DateLiteral) endLiteral;
LocalDate startDate = LocalDate.of(
(int) startInclusiveDate.getYear(),
(int) startInclusiveDate.getMonth(),
(int) startInclusiveDate.getDay()
);
LocalDate endDate = LocalDate.of(
(int) endLiteralDate.getYear(),
(int) endLiteralDate.getMonth(),
(int) endLiteralDate.getDay()
);
if (endExclusive
&& startDate.atStartOfDay(ZoneOffset.UTC).toInstant().toEpochMilli() + ONE_DAY_MILLIS_SECOND
>= endDate.atStartOfDay(ZoneOffset.UTC).toInstant().toEpochMilli()) {
return new SingletonIterator(startInclusive);
}
return new DateLikeRangePartitionValueIterator<>(startDate, endDate, endExclusive,
date -> new DateLiteral(DateFormatUtils.format(date, DateLiteral.JAVA_DATE_FORMAT)));
date -> new DateLiteral(date.getYear(), date.getMonthValue(), date.getDayOfMonth()));
} else if (dataType.isDateV2Type()) {
Date startDate = DateUtils.parseDate(startInclusive.toString(), DateLiteral.JAVA_DATE_FORMAT);
Date endDate = DateUtils.parseDate(endLiteral.toString(), DateLiteral.JAVA_DATE_FORMAT);
DateV2Literal startInclusiveDate = (DateV2Literal) startInclusive;
DateV2Literal endLiteralDate = (DateV2Literal) endLiteral;
LocalDate startDate = LocalDate.of(
(int) startInclusiveDate.getYear(),
(int) startInclusiveDate.getMonth(),
(int) startInclusiveDate.getDay()
);
LocalDate endDate = LocalDate.of(
(int) endLiteralDate.getYear(),
(int) endLiteralDate.getMonth(),
(int) endLiteralDate.getDay()
);
if (endExclusive
&& startDate.atStartOfDay(ZoneOffset.UTC).toInstant().toEpochMilli() + ONE_DAY_MILLIS_SECOND
>= endDate.atStartOfDay(ZoneOffset.UTC).toInstant().toEpochMilli()) {
return new SingletonIterator(startInclusive);
}
return new DateLikeRangePartitionValueIterator<>(startDate, endDate, endExclusive,
date -> new DateV2Literal(DateFormatUtils.format(date, DateLiteral.JAVA_DATE_FORMAT)));
date -> new DateV2Literal(date.getYear(), date.getMonthValue(), date.getDayOfMonth()));
}
// unsupported type
return Iterators.singletonIterator(slot);
@ -231,16 +297,16 @@ public class PartitionRangeExpander {
}
private class DateLikeRangePartitionValueIterator<L extends Literal>
extends RangePartitionValueIterator<Date, L> {
extends RangePartitionValueIterator<LocalDate, L> {
public DateLikeRangePartitionValueIterator(
Date startInclusive, Date finish, boolean endExclusive, Function<Date, L> toLiteral) {
LocalDate startInclusive, LocalDate finish, boolean endExclusive, Function<LocalDate, L> toLiteral) {
super(startInclusive, finish, endExclusive, toLiteral);
}
@Override
protected Date doGetNext(Date current) {
return DateUtils.addDays(current, 1);
protected LocalDate doGetNext(LocalDate current) {
return current.plusDays(1);
}
}

View File

@ -32,7 +32,7 @@ import org.apache.doris.nereids.trees.plans.logical.LogicalOlapScan;
import org.apache.doris.nereids.util.Utils;
import org.apache.doris.qe.ConnectContext;
import com.google.common.collect.ImmutableList;
import com.google.common.collect.Maps;
import java.util.ArrayList;
import java.util.List;
@ -58,13 +58,16 @@ public class PruneOlapScanPartition extends OneRewriteRuleFactory {
return filter;
}
Map<String, Slot> scanOutput = scan.getOutput()
.stream()
.collect(Collectors.toMap(slot -> slot.getName().toLowerCase(), Function.identity()));
List<Slot> output = scan.getOutput();
Map<String, Slot> scanOutput = Maps.newHashMapWithExpectedSize(output.size() * 2);
for (Slot slot : output) {
scanOutput.put(slot.getName().toLowerCase(), slot);
}
PartitionInfo partitionInfo = table.getPartitionInfo();
List<Slot> partitionSlots = new ArrayList<>();
for (Column column : partitionInfo.getPartitionColumns()) {
List<Column> partitionColumns = partitionInfo.getPartitionColumns();
List<Slot> partitionSlots = new ArrayList<>(partitionColumns.size());
for (Column column : partitionColumns) {
Slot slot = scanOutput.get(column.getName().toLowerCase());
if (slot == null) {
return filter;
@ -84,16 +87,16 @@ public class PruneOlapScanPartition extends OneRewriteRuleFactory {
.filter(id -> manuallySpecifiedPartitions.contains(id))
.collect(Collectors.toMap(Function.identity(), id -> allPartitions.get(id)));
}
List<Long> prunedPartitions = new ArrayList<>(PartitionPruner.prune(
List<Long> prunedPartitions = PartitionPruner.prune(
partitionSlots, filter.getPredicate(), idToPartitions, ctx.cascadesContext,
PartitionTableType.OLAP));
PartitionTableType.OLAP);
if (prunedPartitions.isEmpty()) {
return new LogicalEmptyRelation(
ConnectContext.get().getStatementContext().getNextRelationId(),
filter.getOutput());
}
LogicalOlapScan rewrittenScan = scan.withSelectedPartitionIds(ImmutableList.copyOf(prunedPartitions));
LogicalOlapScan rewrittenScan = scan.withSelectedPartitionIds(prunedPartitions);
return new LogicalFilter<>(filter.getConjuncts(), rewrittenScan);
}).toRule(RuleType.OLAP_SCAN_PARTITION_PRUNE);
}

View File

@ -294,51 +294,52 @@ public abstract class Literal extends Expression implements LeafExpression, Comp
} else if (literalExpr instanceof org.apache.doris.analysis.NullLiteral) {
return new NullLiteral(dataType);
}
String stringValue = literalExpr.getStringValue();
if (dataType.isBooleanType()) {
return ((BoolLiteral) literalExpr).getValue() ? BooleanLiteral.TRUE : BooleanLiteral.FALSE;
} else if (dataType.isTinyIntType()) {
return new TinyIntLiteral(Byte.parseByte(stringValue));
} else if (dataType.isSmallIntType()) {
return new SmallIntLiteral(Short.parseShort(stringValue));
} else if (dataType.isIntegerType()) {
return new IntegerLiteral(Integer.parseInt(stringValue));
} else if (dataType.isBigIntType()) {
return new BigIntLiteral(Long.parseLong(stringValue));
} else if (dataType.isLargeIntType()) {
return new LargeIntLiteral(new BigInteger(stringValue));
} else if (dataType.isStringType()) {
return new StringLiteral(stringValue);
} else if (dataType.isCharType()) {
return new CharLiteral(stringValue, ((CharType) dataType).getLen());
} else if (dataType.isVarcharType()) {
return new VarcharLiteral(stringValue, ((VarcharType) dataType).getLen());
} else if (dataType.isFloatType()) {
return new FloatLiteral(Float.parseFloat(stringValue));
} else if (dataType.isDoubleType()) {
return new DoubleLiteral(Double.parseDouble(stringValue));
} else if (dataType.isDecimalV2Type()) {
return new DecimalLiteral((DecimalV2Type) dataType, new BigDecimal(stringValue));
} else if (dataType.isDecimalV3Type()) {
return new DecimalV3Literal((DecimalV3Type) dataType, new BigDecimal(stringValue));
} else if (dataType.isDateType()) {
return new DateLiteral(stringValue);
} else if (dataType.isDateV2Type()) {
return new DateV2Literal(stringValue);
} else if (dataType.isDateTimeType()) {
return new DateTimeLiteral(stringValue);
} else if (dataType.isDateTimeV2Type()) {
return new DateTimeV2Literal(stringValue);
} else if (dataType.isJsonType()) {
return new JsonLiteral(stringValue);
} else if (dataType.isIPv4Type()) {
return new IPv4Literal(stringValue);
} else if (dataType.isIPv6Type()) {
return new IPv6Literal(stringValue);
} else {
throw new AnalysisException("Unsupported convert the " + literalExpr.getType()
+ " of legacy literal to nereids literal");
// fast path
switch (type.getPrimitiveType()) {
case DATEV2: {
org.apache.doris.analysis.DateLiteral dateLiteral = (org.apache.doris.analysis.DateLiteral) literalExpr;
return new DateV2Literal(dateLiteral.getYear(), dateLiteral.getMonth(), dateLiteral.getDay());
}
case DATE: {
org.apache.doris.analysis.DateLiteral dateLiteral = (org.apache.doris.analysis.DateLiteral) literalExpr;
return new DateLiteral(dateLiteral.getYear(), dateLiteral.getMonth(), dateLiteral.getDay());
}
case BOOLEAN: {
return ((BoolLiteral) literalExpr).getValue() ? BooleanLiteral.TRUE : BooleanLiteral.FALSE;
}
default: {
}
}
// slow path
String stringValue = literalExpr.getStringValue();
switch (type.getPrimitiveType()) {
case TINYINT: return new TinyIntLiteral(Byte.parseByte(stringValue));
case SMALLINT: return new SmallIntLiteral(Short.parseShort(stringValue));
case INT: return new IntegerLiteral(Integer.parseInt(stringValue));
case BIGINT: return new BigIntLiteral(Long.parseLong(stringValue));
case LARGEINT: return new LargeIntLiteral(new BigInteger(stringValue));
case STRING: return new StringLiteral(stringValue);
case CHAR: return new CharLiteral(stringValue, ((CharType) dataType).getLen());
case VARCHAR: return new VarcharLiteral(stringValue, ((VarcharType) dataType).getLen());
case FLOAT: return new FloatLiteral(Float.parseFloat(stringValue));
case DOUBLE: return new DoubleLiteral(Double.parseDouble(stringValue));
case DECIMALV2: return new DecimalLiteral((DecimalV2Type) dataType, new BigDecimal(stringValue));
case DECIMAL32:
case DECIMAL64:
case DECIMAL128:
case DECIMAL256: {
return new DecimalV3Literal((DecimalV3Type) dataType, new BigDecimal(stringValue));
}
case DATETIME: return new DateTimeLiteral(stringValue);
case DATETIMEV2: return new DateTimeV2Literal(stringValue);
case JSONB: return new JsonLiteral(stringValue);
case IPV4: return new IPv4Literal(stringValue);
case IPV6: return new IPv6Literal(stringValue);
default: {
}
}
throw new AnalysisException("Unsupported convert the " + literalExpr.getType()
+ " of legacy literal to nereids literal");
}
@Override

View File

@ -325,63 +325,58 @@ public abstract class DataType {
*/
@Developing // should support map, struct
public static DataType fromCatalogType(Type type) {
if (type.isBoolean()) {
return BooleanType.INSTANCE;
} else if (type.getPrimitiveType() == Type.TINYINT.getPrimitiveType()) {
return TinyIntType.INSTANCE;
} else if (type.getPrimitiveType() == Type.SMALLINT.getPrimitiveType()) {
return SmallIntType.INSTANCE;
} else if (type.getPrimitiveType() == Type.INT.getPrimitiveType()) {
return IntegerType.INSTANCE;
} else if (type.getPrimitiveType() == Type.BIGINT.getPrimitiveType()) {
return BigIntType.INSTANCE;
} else if (type.getPrimitiveType() == Type.LARGEINT.getPrimitiveType()) {
return LargeIntType.INSTANCE;
} else if (type.getPrimitiveType() == Type.FLOAT.getPrimitiveType()) {
return FloatType.INSTANCE;
} else if (type.getPrimitiveType() == Type.DOUBLE.getPrimitiveType()) {
return DoubleType.INSTANCE;
} else if (type.isNull()) {
return NullType.INSTANCE;
} else if (type.isDatetimeV2()) {
return DateTimeV2Type.of(((ScalarType) type).getScalarScale());
} else if (type.isDatetime()) {
return DateTimeType.INSTANCE;
} else if (type.isDateV2()) {
return DateV2Type.INSTANCE;
} else if (type.isDateType()) {
return DateType.INSTANCE;
} else if (type.isTimeV2()) {
return TimeV2Type.INSTANCE;
} else if (type.isTime()) {
return TimeType.INSTANCE;
} else if (type.isHllType()) {
return HllType.INSTANCE;
} else if (type.isBitmapType()) {
return BitmapType.INSTANCE;
} else if (type.isQuantileStateType()) {
return QuantileStateType.INSTANCE;
} else if (type.getPrimitiveType() == org.apache.doris.catalog.PrimitiveType.CHAR) {
return CharType.createCharType(type.getLength());
} else if (type.getPrimitiveType() == org.apache.doris.catalog.PrimitiveType.VARCHAR) {
return VarcharType.createVarcharType(type.getLength());
} else if (type.getPrimitiveType() == org.apache.doris.catalog.PrimitiveType.STRING) {
return StringType.INSTANCE;
} else if (type.isDecimalV3()) {
ScalarType scalarType = (ScalarType) type;
int precision = scalarType.getScalarPrecision();
int scale = scalarType.getScalarScale();
return DecimalV3Type.createDecimalV3TypeNoCheck(precision, scale);
} else if (type.isDecimalV2()) {
ScalarType scalarType = (ScalarType) type;
int precision = scalarType.getScalarPrecision();
int scale = scalarType.getScalarScale();
return DecimalV2Type.createDecimalV2Type(precision, scale);
} else if (type.isJsonbType()) {
return JsonType.INSTANCE;
} else if (type.isVariantType()) {
return VariantType.INSTANCE;
} else if (type.isStructType()) {
switch (type.getPrimitiveType()) {
case BOOLEAN: return BooleanType.INSTANCE;
case TINYINT: return TinyIntType.INSTANCE;
case SMALLINT: return SmallIntType.INSTANCE;
case INT: return IntegerType.INSTANCE;
case BIGINT: return BigIntType.INSTANCE;
case LARGEINT: return LargeIntType.INSTANCE;
case FLOAT: return FloatType.INSTANCE;
case DOUBLE: return DoubleType.INSTANCE;
case NULL_TYPE: return NullType.INSTANCE;
case DATETIMEV2: return DateTimeV2Type.of(((ScalarType) type).getScalarScale());
case DATETIME: return DateTimeType.INSTANCE;
case DATEV2: return DateV2Type.INSTANCE;
case DATE: return DateType.INSTANCE;
case TIMEV2: return TimeV2Type.INSTANCE;
case TIME: return TimeType.INSTANCE;
case HLL: return HllType.INSTANCE;
case BITMAP: return BitmapType.INSTANCE;
case QUANTILE_STATE: return QuantileStateType.INSTANCE;
case CHAR: return CharType.createCharType(type.getLength());
case VARCHAR: return VarcharType.createVarcharType(type.getLength());
case STRING: return StringType.INSTANCE;
case VARIANT: return VariantType.INSTANCE;
case JSONB: return JsonType.INSTANCE;
case IPV4: return IPv4Type.INSTANCE;
case IPV6: return IPv6Type.INSTANCE;
case AGG_STATE: {
org.apache.doris.catalog.AggStateType catalogType = ((org.apache.doris.catalog.AggStateType) type);
List<DataType> types = catalogType.getSubTypes().stream().map(DataType::fromCatalogType)
.collect(Collectors.toList());
return new AggStateType(catalogType.getFunctionName(), types, catalogType.getSubTypeNullables());
}
case DECIMALV2: {
ScalarType scalarType = (ScalarType) type;
int precision = scalarType.getScalarPrecision();
int scale = scalarType.getScalarScale();
return DecimalV2Type.createDecimalV2Type(precision, scale);
}
case DECIMAL32:
case DECIMAL64:
case DECIMAL128:
case DECIMAL256: {
ScalarType scalarType = (ScalarType) type;
int precision = scalarType.getScalarPrecision();
int scale = scalarType.getScalarScale();
return DecimalV3Type.createDecimalV3TypeNoCheck(precision, scale);
}
default: {
}
}
if (type.isStructType()) {
List<StructField> structFields = ((org.apache.doris.catalog.StructType) (type)).getFields().stream()
.map(cf -> new StructField(cf.getName(), fromCatalogType(cf.getType()),
cf.getContainsNull(), cf.getComment() == null ? "" : cf.getComment()))
@ -393,15 +388,6 @@ public abstract class DataType {
} else if (type.isArrayType()) {
org.apache.doris.catalog.ArrayType arrayType = (org.apache.doris.catalog.ArrayType) type;
return ArrayType.of(fromCatalogType(arrayType.getItemType()), arrayType.getContainsNull());
} else if (type.isAggStateType()) {
org.apache.doris.catalog.AggStateType catalogType = ((org.apache.doris.catalog.AggStateType) type);
List<DataType> types = catalogType.getSubTypes().stream().map(DataType::fromCatalogType)
.collect(Collectors.toList());
return new AggStateType(catalogType.getFunctionName(), types, catalogType.getSubTypeNullables());
} else if (type.isIPv4()) {
return IPv4Type.INSTANCE;
} else if (type.isIPv6()) {
return IPv6Type.INSTANCE;
} else {
return UnsupportedType.INSTANCE;
}

View File

@ -224,6 +224,22 @@ public class Utils {
/** allCombinations */
public static <T> List<List<T>> allCombinations(List<List<T>> lists) {
if (lists.size() == 1) {
List<T> first = lists.get(0);
if (first.size() == 1) {
return lists;
}
List<List<T>> result = Lists.newArrayListWithCapacity(lists.size());
for (T item : first) {
result.add(ImmutableList.of(item));
}
return result;
} else {
return doAllCombinations(lists);
}
}
private static <T> List<List<T>> doAllCombinations(List<List<T>> lists) {
int size = lists.size();
if (size == 0) {
return ImmutableList.of();