DruidDateTimeUtils.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.calcite.adapter.druid;
import org.apache.calcite.avatica.util.TimeUnitRange;
import org.apache.calcite.rel.type.RelDataType;
import org.apache.calcite.rex.RexCall;
import org.apache.calcite.rex.RexInputRef;
import org.apache.calcite.rex.RexLiteral;
import org.apache.calcite.rex.RexNode;
import org.apache.calcite.sql.SqlKind;
import org.apache.calcite.sql.type.SqlTypeName;
import org.apache.calcite.util.DateString;
import org.apache.calcite.util.RangeSets;
import org.apache.calcite.util.Sarg;
import org.apache.calcite.util.TimestampString;
import org.apache.calcite.util.Util;
import org.apache.calcite.util.trace.CalciteTrace;
import com.google.common.collect.BoundType;
import com.google.common.collect.ImmutableList;
import com.google.common.collect.ImmutableRangeSet;
import com.google.common.collect.Range;
import com.google.common.collect.TreeRangeSet;
import org.checkerframework.checker.nullness.qual.Nullable;
import org.joda.time.Interval;
import org.joda.time.Period;
import org.joda.time.chrono.ISOChronology;
import org.slf4j.Logger;
import java.util.ArrayList;
import java.util.List;
import static java.util.Objects.requireNonNull;
/**
* Utilities for generating intervals from RexNode.
*/
@SuppressWarnings({"rawtypes", "unchecked" })
public class DruidDateTimeUtils {
protected static final Logger LOGGER = CalciteTrace.getPlannerTracer();
private DruidDateTimeUtils() {
}
/**
* Generates a list of {@link Interval}s equivalent to a given
* expression. Assumes that all the predicates in the input
* reference a single column: the timestamp column.
*/
public static @Nullable List<Interval> createInterval(RexNode e) {
final List<Range<Long>> ranges = extractRanges(e, false);
if (ranges == null) {
// We did not succeed, bail out
return null;
}
final TreeRangeSet condensedRanges = TreeRangeSet.create();
for (Range r : ranges) {
condensedRanges.add(r);
}
LOGGER.debug("Inferred ranges on interval : {}", condensedRanges);
return toInterval(ImmutableList.<Range<Long>>copyOf(condensedRanges.asRanges()));
}
protected static List<Interval> toInterval(
List<Range<Long>> ranges) {
List<Interval> intervals = Util.transform(ranges, range -> {
if (!range.hasLowerBound() && !range.hasUpperBound()) {
return DruidTable.DEFAULT_INTERVAL;
}
long start = range.hasLowerBound()
? range.lowerEndpoint().longValue()
: DruidTable.DEFAULT_INTERVAL.getStartMillis();
long end = range.hasUpperBound()
? range.upperEndpoint().longValue()
: DruidTable.DEFAULT_INTERVAL.getEndMillis();
if (range.hasLowerBound()
&& range.lowerBoundType() == BoundType.OPEN) {
start++;
}
if (range.hasUpperBound()
&& range.upperBoundType() == BoundType.CLOSED) {
end++;
}
return new Interval(start, end, ISOChronology.getInstanceUTC());
});
if (LOGGER.isDebugEnabled()) {
LOGGER.debug("Converted time ranges " + ranges + " to interval " + intervals);
}
return intervals;
}
protected static @Nullable List<Range<Long>> extractRanges(RexNode node, boolean withNot) {
switch (node.getKind()) {
case EQUALS:
case LESS_THAN:
case LESS_THAN_OR_EQUAL:
case GREATER_THAN:
case GREATER_THAN_OR_EQUAL:
case DRUID_IN:
case SEARCH:
return leafToRanges((RexCall) node, withNot);
case NOT:
return extractRanges(((RexCall) node).getOperands().get(0), !withNot);
case OR: {
RexCall call = (RexCall) node;
List<Range<Long>> intervals = new ArrayList<>();
for (RexNode child : call.getOperands()) {
List<Range<Long>> extracted =
extractRanges(child, withNot);
if (extracted != null) {
intervals.addAll(extracted);
}
}
return intervals;
}
case AND: {
RexCall call = (RexCall) node;
List<Range<Long>> ranges = new ArrayList<>();
for (RexNode child : call.getOperands()) {
List<Range<Long>> extractedRanges =
extractRanges(child, false);
if (extractedRanges == null || extractedRanges.isEmpty()) {
// We could not extract, we bail out
return null;
}
if (ranges.isEmpty()) {
ranges.addAll(extractedRanges);
continue;
}
List<Range<Long>> overlapped = new ArrayList<>();
for (Range current : ranges) {
for (Range interval : extractedRanges) {
if (current.isConnected(interval)) {
overlapped.add(current.intersection(interval));
}
}
}
ranges = overlapped;
}
return ranges;
}
default:
return null;
}
}
protected static @Nullable List<Range<Long>> leafToRanges(RexCall call, boolean withNot) {
final ImmutableList.Builder<Range<Long>> ranges;
switch (call.getKind()) {
case EQUALS:
case LESS_THAN:
case LESS_THAN_OR_EQUAL:
case GREATER_THAN:
case GREATER_THAN_OR_EQUAL: {
final Long value;
SqlKind kind = call.getKind();
if (call.getOperands().get(0) instanceof RexInputRef
&& literalValue(call.getOperands().get(1)) != null) {
value = requireNonNull(literalValue(call.getOperands().get(1)));
} else if (call.getOperands().get(1) instanceof RexInputRef
&& literalValue(call.getOperands().get(0)) != null) {
value = requireNonNull(literalValue(call.getOperands().get(0)));
kind = kind.reverse();
} else {
return null;
}
switch (kind) {
case LESS_THAN:
return ImmutableList.of(withNot ? Range.atLeast(value) : Range.lessThan(value));
case LESS_THAN_OR_EQUAL:
return ImmutableList.of(withNot ? Range.greaterThan(value) : Range.atMost(value));
case GREATER_THAN:
return ImmutableList.of(withNot ? Range.atMost(value) : Range.greaterThan(value));
case GREATER_THAN_OR_EQUAL:
return ImmutableList.of(withNot ? Range.lessThan(value) : Range.atLeast(value));
default:
if (!withNot) {
return ImmutableList.of(Range.closed(value, value));
}
return ImmutableList.of(Range.lessThan(value), Range.greaterThan(value));
}
}
case BETWEEN: {
final Long value1;
final Long value2;
if (literalValue(call.getOperands().get(2)) != null
&& literalValue(call.getOperands().get(3)) != null) {
value1 = requireNonNull(literalValue(call.getOperands().get(2)));
value2 = requireNonNull(literalValue(call.getOperands().get(3)));
} else {
return null;
}
boolean inverted = value1.compareTo(value2) > 0;
if (!withNot) {
return ImmutableList.of(
inverted ? Range.closed(value2, value1) : Range.closed(value1, value2));
}
return ImmutableList.of(Range.lessThan(inverted ? value2 : value1),
Range.greaterThan(inverted ? value1 : value2));
}
case DRUID_IN:
ranges = ImmutableList.builder();
for (RexNode operand : Util.skip(call.operands)) {
final Long element = literalValue(operand);
if (element == null) {
return null;
}
if (withNot) {
ranges.add(Range.lessThan(element));
ranges.add(Range.greaterThan(element));
} else {
ranges.add(Range.closed(element, element));
}
}
return ranges.build();
case SEARCH:
final RexLiteral right = (RexLiteral) call.operands.get(1);
final Sarg<?> sarg = requireNonNull(right.getValueAs(Sarg.class));
ranges = ImmutableList.builder();
for (Range range : sarg.rangeSet.asRanges()) {
Range<Long> range2 = RangeSets.copy(range, DruidDateTimeUtils::toLong);
if (withNot) {
ranges.addAll(ImmutableRangeSet.of(range2).complement().asRanges());
} else {
ranges.add(range2);
}
}
return ranges.build();
default:
return null;
}
}
private static Long toLong(Comparable comparable) {
if (comparable instanceof TimestampString) {
TimestampString timestampString = (TimestampString) comparable;
return timestampString.getMillisSinceEpoch();
}
if (comparable instanceof DateString) {
DateString dataString = (DateString) comparable;
return dataString.getMillisSinceEpoch();
}
throw new AssertionError("unsupported type: " + comparable.getClass());
}
/**
* Returns the literal value for the given node, assuming it is a literal with
* datetime type, or a cast that only alters nullability on top of a literal with
* datetime type.
*/
protected static @Nullable Long literalValue(RexNode node) {
switch (node.getKind()) {
case LITERAL:
switch (((RexLiteral) node).getTypeName()) {
case TIMESTAMP:
case TIMESTAMP_WITH_LOCAL_TIME_ZONE:
TimestampString tsVal = ((RexLiteral) node).getValueAs(TimestampString.class);
if (tsVal == null) {
return null;
}
return tsVal.getMillisSinceEpoch();
case DATE:
DateString dateVal = ((RexLiteral) node).getValueAs(DateString.class);
if (dateVal == null) {
return null;
}
return dateVal.getMillisSinceEpoch();
default:
break;
}
break;
case CAST:
// Normally all CASTs are eliminated by now by constant reduction.
// But when HiveExecutor is used there may be a cast that changes only
// nullability, from TIMESTAMP NOT NULL literal to TIMESTAMP literal.
// We can handle that case by traversing the dummy CAST.
assert node instanceof RexCall;
final RexCall call = (RexCall) node;
final RexNode operand = call.getOperands().get(0);
final RelDataType callType = call.getType();
final RelDataType operandType = operand.getType();
if (operand.getKind() == SqlKind.LITERAL
&& callType.getSqlTypeName() == operandType.getSqlTypeName()
&& (callType.getSqlTypeName() == SqlTypeName.DATE
|| callType.getSqlTypeName() == SqlTypeName.TIMESTAMP
|| callType.getSqlTypeName() == SqlTypeName.TIMESTAMP_WITH_LOCAL_TIME_ZONE)
&& callType.isNullable()
&& !operandType.isNullable()) {
return literalValue(operand);
}
break;
default:
break;
}
return null;
}
/**
* Infers granularity from a time unit.
* It supports {@code FLOOR(<time> TO <timeunit>)}
* and {@code EXTRACT(<timeunit> FROM <time>)}.
* Returns null if it cannot be inferred.
*
* @param node the Rex node
* @return the granularity, or null if it cannot be inferred
*/
public static @Nullable Granularity extractGranularity(RexNode node, String timeZone) {
final int valueIndex;
final int flagIndex;
if (TimeExtractionFunction.isValidTimeExtract(node)) {
flagIndex = 0;
valueIndex = 1;
} else if (TimeExtractionFunction.isValidTimeFloor(node)) {
valueIndex = 0;
flagIndex = 1;
} else {
// We can only infer granularity from floor and extract.
return null;
}
final RexCall call = (RexCall) node;
final RexNode value = call.operands.get(valueIndex);
final RexLiteral flag = (RexLiteral) call.operands.get(flagIndex);
final TimeUnitRange timeUnit =
requireNonNull((TimeUnitRange) flag.getValue());
final RelDataType valueType = value.getType();
if (valueType.getSqlTypeName() == SqlTypeName.DATE
|| valueType.getSqlTypeName() == SqlTypeName.TIMESTAMP) {
// We use 'UTC' for date/timestamp type as Druid needs timezone information
return Granularities.createGranularity(timeUnit, "UTC");
} else if (valueType.getSqlTypeName() == SqlTypeName.TIMESTAMP_WITH_LOCAL_TIME_ZONE) {
return Granularities.createGranularity(timeUnit, timeZone);
}
// Type not recognized
return null;
}
/**
* Converts a granularity to ISO period format.
*
* @param type Druid Granularity to translate as period of time
*
* @return String representing the granularity as ISO8601 Period of Time; null
* for unknown case
*/
public static @Nullable String toISOPeriodFormat(Granularity.Type type) {
switch (type) {
case SECOND:
return Period.seconds(1).toString();
case MINUTE:
return Period.minutes(1).toString();
case HOUR:
return Period.hours(1).toString();
case DAY:
return Period.days(1).toString();
case WEEK:
return Period.weeks(1).toString();
case MONTH:
return Period.months(1).toString();
case QUARTER:
return Period.months(3).toString();
case YEAR:
return Period.years(1).toString();
default:
return null;
}
}
/**
* Translates a Calcite {@link TimeUnitRange} to a Druid {@link Granularity}.
*
* @param timeUnit Calcite Time unit to convert
*
* @return Druid Granularity or null
*/
public static Granularity.@Nullable Type toDruidGranularity(TimeUnitRange timeUnit) {
if (timeUnit == null) {
return null;
}
switch (timeUnit) {
case YEAR:
return Granularity.Type.YEAR;
case QUARTER:
return Granularity.Type.QUARTER;
case MONTH:
return Granularity.Type.MONTH;
case WEEK:
return Granularity.Type.WEEK;
case DAY:
return Granularity.Type.DAY;
case HOUR:
return Granularity.Type.HOUR;
case MINUTE:
return Granularity.Type.MINUTE;
case SECOND:
return Granularity.Type.SECOND;
default:
return null;
}
}
}