Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
3 changes: 3 additions & 0 deletions native/core/src/execution/serde.rs
Original file line number Diff line number Diff line change
Expand Up @@ -102,6 +102,9 @@ pub fn to_arrow_datatype(dt_value: &DataType) -> ArrowDataType {
// Spark's DayTimeIntervalType stores microseconds in an int64, which matches Arrow
// Duration(Microsecond) rather than the lossy Interval(DayTime) {days, millis} layout.
DataTypeId::DayTimeInterval => ArrowDataType::Duration(TimeUnit::Microsecond),
// Spark's CalendarIntervalType stores months, days, and microseconds. Arrow stores the
// same components with nanosecond precision.
DataTypeId::CalendarInterval => ArrowDataType::Interval(IntervalUnit::MonthDayNano),
DataTypeId::Null => ArrowDataType::Null,
DataTypeId::List => match dt_value
.type_info
Expand Down
1 change: 1 addition & 0 deletions native/proto/src/proto/types.proto
Original file line number Diff line number Diff line change
Expand Up @@ -62,6 +62,7 @@ message DataType {
TIME = 17;
YEAR_MONTH_INTERVAL = 18;
DAY_TIME_INTERVAL = 19;
CALENDAR_INTERVAL = 20;
}
DataTypeId type_id = 1;

Expand Down
75 changes: 75 additions & 0 deletions spark/src/main/java/org/apache/comet/vector/CometPlainVector.java
Original file line number Diff line number Diff line change
Expand Up @@ -27,6 +27,8 @@
import org.apache.arrow.vector.*;
import org.apache.arrow.vector.util.TransferPair;
import org.apache.parquet.Preconditions;
import org.apache.spark.sql.types.DataTypes;
import org.apache.spark.sql.vectorized.ColumnVector;
import org.apache.spark.unsafe.Platform;
import org.apache.spark.unsafe.types.UTF8String;

Expand All @@ -41,6 +43,7 @@ public class CometPlainVector extends CometDecodedVector {
private final long valueBufferAddress;
private final long offsetBufferAddress;
private final boolean isBaseFixedWidthVector;
private final ColumnVector[] intervalChildren;
// True when the variable-width offsets are 64-bit (LargeVarChar / LargeVarBinary) rather than
// the usual 32-bit. PyArrow UDFs can hand back large_string / large_binary columns.
private final boolean isLargeVarWidth;
Expand Down Expand Up @@ -73,6 +76,16 @@ public CometPlainVector(ValueVector vector, boolean isUuid) {
this.offsetBufferAddress = -1;
this.isLargeVarWidth = false;
}
if (vector instanceof IntervalMonthDayNanoVector) {
this.intervalChildren =
new ColumnVector[] {
new IntervalChildVector(this, 0),
new IntervalChildVector(this, 1),
new IntervalChildVector(this, 2)
};
} else {
this.intervalChildren = null;
}
}

@Override
Expand Down Expand Up @@ -178,6 +191,14 @@ public byte[] getBinary(int rowId) {
return result;
}

@Override
public ColumnVector getChild(int ordinal) {
if (intervalChildren == null || ordinal < 0 || ordinal >= intervalChildren.length) {
return super.getChild(ordinal);
}
return intervalChildren[ordinal];
}

@Override
public CDataDictionaryProvider getDictionaryProvider() {
return null;
Expand All @@ -204,4 +225,58 @@ private static UUID convertToUuid(byte[] buf) {
long leastSigBits = bb.getLong();
return new UUID(mostSigBits, leastSigBits);
}

private static final class IntervalChildVector extends CometVector {
private final CometPlainVector parent;
private final int ordinal;

private IntervalChildVector(CometPlainVector parent, int ordinal) {
super(ordinal == 2 ? DataTypes.LongType : DataTypes.IntegerType);
this.parent = parent;
this.ordinal = ordinal;
}

@Override
public int getInt(int rowId) {
return Platform.getInt(null, parent.valueBufferAddress + rowId * 16L + ordinal * 4L);
}

@Override
public long getLong(int rowId) {
return Platform.getLong(null, parent.valueBufferAddress + rowId * 16L + 8L) / 1000L;
}

@Override
public boolean isNullAt(int rowId) {
return parent.isNullAt(rowId);
}

@Override
public boolean hasNull() {
return parent.hasNull();
}

@Override
public int numNulls() {
return parent.numNulls();
}

@Override
public int numValues() {
return parent.numValues();
}

@Override
public ValueVector getValueVector() {
return parent.getValueVector();
}

@Override
public CometVector slice(int offset, int length) {
throw new UnsupportedOperationException("Interval child vectors cannot be sliced");
}

@Override
public void close() {}
}
}
4 changes: 2 additions & 2 deletions spark/src/main/scala/org/apache/comet/DataTypeSupport.scala
Original file line number Diff line number Diff line change
Expand Up @@ -50,8 +50,8 @@ trait DataTypeSupport {

dt match {
case BooleanType | ByteType | ShortType | IntegerType | LongType | FloatType | DoubleType |
BinaryType | StringType | _: DecimalType | DateType | TimestampType |
TimestampNTZType =>
BinaryType | StringType | _: DecimalType | DateType | TimestampType | TimestampNTZType |
CalendarIntervalType =>
true
case StructType(fields) =>
fields.nonEmpty && fields.forall(f =>
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -73,6 +73,7 @@ object CometBatchKernelCodegen extends Logging with CometExprTraitShim {
case "VarBinaryVector" => classOf[VarBinaryVector]
case "IntervalYearVector" => classOf[IntervalYearVector]
case "DurationVector" => classOf[DurationVector]
case "IntervalMonthDayNanoVector" => classOf[IntervalMonthDayNanoVector]
case other => throw new IllegalArgumentException(s"unknown Arrow vector class: $other")
}

Expand All @@ -86,7 +87,7 @@ object CometBatchKernelCodegen extends Logging with CometExprTraitShim {
case _: DecimalType => true
case _: StringType | _: BinaryType => true
case DateType | TimestampType | TimestampNTZType => true
case _: YearMonthIntervalType | _: DayTimeIntervalType => true
case _: YearMonthIntervalType | _: DayTimeIntervalType | CalendarIntervalType => true
case ArrayType(inner, _) => isSupportedDataType(inner)
case st: StructType => st.fields.forall(f => isSupportedDataType(f.dataType))
case mt: MapType => isSupportedDataType(mt.keyType) && isSupportedDataType(mt.valueType)
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -62,7 +62,8 @@ private[codegen] object CometBatchKernelCodegenInput {
classOf[Float8Vector],
classOf[DateDayVector],
classOf[TimeStampMicroVector],
classOf[TimeStampMicroTZVector])
classOf[TimeStampMicroTZVector],
classOf[IntervalMonthDayNanoVector])
private val cometPlainVectorName: String = classOf[CometPlainVector].getName

/** Emit kernel typed-vector field declarations for every level of every input column. */
Expand Down Expand Up @@ -137,6 +138,10 @@ private[codegen] object CometBatchKernelCodegenInput {
cls == classOf[TimeStampMicroTZVector] =>
s" case $ord: return this.col$ord.getLong(this.rowIdx);"
}
val intervalCases = withOrd.collect {
case (ArrowColumnSpec(cls, _), ord) if cls == classOf[IntervalMonthDayNanoVector] =>
s" case $ord: return this.col$ord.getInterval(this.rowIdx);"
}
val floatCases = withOrd.collect {
case (ArrowColumnSpec(cls, _), ord) if cls == classOf[Float4Vector] =>
s" case $ord: return this.col$ord.getFloat(this.rowIdx);"
Expand Down Expand Up @@ -196,6 +201,10 @@ private[codegen] object CometBatchKernelCodegenInput {
emitOrdinalSwitch("public long getLong(int ordinal)", "getLong", longCases),
emitOrdinalSwitch("public float getFloat(int ordinal)", "getFloat", floatCases),
emitOrdinalSwitch("public double getDouble(int ordinal)", "getDouble", doubleCases),
emitOrdinalSwitch(
"public org.apache.spark.unsafe.types.CalendarInterval getInterval(int ordinal)",
"getInterval",
intervalCases),
emitOrdinalSwitch(
"public org.apache.spark.sql.types.Decimal getDecimal(" +
"int ordinal, int precision, int scale)",
Expand Down Expand Up @@ -592,6 +601,7 @@ private[codegen] object CometBatchKernelCodegenInput {
case ShortType => s"getShort($idx)"
case IntegerType | DateType => s"getInt($idx)"
case LongType | TimestampType | TimestampNTZType => s"getLong($idx)"
case CalendarIntervalType => s"getInterval($idx)"
case FloatType => s"getFloat($idx)"
case DoubleType => s"getDouble($idx)"
case d: DecimalType => s"getDecimal($idx, ${d.precision}, ${d.scale})"
Expand Down Expand Up @@ -697,6 +707,11 @@ private[codegen] object CometBatchKernelCodegenInput {
| public long getLong(int i) {
| return $childField.getLong(startIndex + i);
| }""".stripMargin
case CalendarIntervalType =>
s""" @Override
| public org.apache.spark.unsafe.types.CalendarInterval getInterval(int i) {
|$nullGuard return $childField.getInterval(startIndex + i);
| }""".stripMargin
case FloatType =>
s""" @Override
| public float getFloat(int i) {
Expand Down Expand Up @@ -847,6 +862,10 @@ private[codegen] object CometBatchKernelCodegenInput {
s" case $fi: return ${path}_f$fi.getInt(this.rowIdx);"
case LongType | TimestampType | TimestampNTZType =>
s" case $fi: return ${path}_f$fi.getLong(this.rowIdx);"
case CalendarIntervalType =>
s""" case $fi: {
|$guard return ${path}_f$fi.getInterval(this.rowIdx);
| }""".stripMargin
case FloatType =>
s" case $fi: return ${path}_f$fi.getFloat(this.rowIdx);"
case DoubleType =>
Expand Down Expand Up @@ -900,6 +919,10 @@ private[codegen] object CometBatchKernelCodegenInput {
f.sparkType == TimestampNTZType =>
fieldReadScalar(fi, LongType, f.nullable)
}
val intervalCases = scalarOrd.collect {
case (f, fi) if f.sparkType == CalendarIntervalType =>
fieldReadScalar(fi, CalendarIntervalType, f.nullable)
}
val floatCases =
scalarOrd.collect {
case (f, fi) if f.sparkType == FloatType =>
Expand Down Expand Up @@ -944,6 +967,10 @@ private[codegen] object CometBatchKernelCodegenInput {
structSwitch("public long getLong(int ordinal)", "getLong", longCases),
structSwitch("public float getFloat(int ordinal)", "getFloat", floatCases),
structSwitch("public double getDouble(int ordinal)", "getDouble", doubleCases),
structSwitch(
"public org.apache.spark.unsafe.types.CalendarInterval getInterval(int ordinal)",
"getInterval",
intervalCases),
structSwitch(
"public org.apache.spark.sql.types.Decimal getDecimal(" +
"int ordinal, int precision, int scale)",
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -173,6 +173,7 @@ private[codegen] object CometBatchKernelCodegenOutput {
case TimestampNTZType => classOf[TimeStampMicroVector].getName
case _: YearMonthIntervalType => classOf[IntervalYearVector].getName
case _: DayTimeIntervalType => classOf[DurationVector].getName
case CalendarIntervalType => classOf[IntervalMonthDayNanoVector].getName
case _: ArrayType => classOf[ListVector].getName
case _: StructType => classOf[StructVector].getName
case _: MapType => classOf[MapVector].getName
Expand Down Expand Up @@ -216,6 +217,14 @@ private[codegen] object CometBatchKernelCodegenOutput {
// DayTimeIntervalType -> DurationVector.set(int, long micros).
val set = if (nested) "setSafe" else "set"
OutputEmit("", s"$targetVec.$set($idx, $source);")
case CalendarIntervalType =>
val set = if (nested) "setSafe" else "set"
val interval = ctx.freshName("interval")
OutputEmit(
"",
s"""org.apache.spark.unsafe.types.CalendarInterval $interval = $source;
|$targetVec.$set($idx, $interval.months, $interval.days,
| java.lang.Math.multiplyExact($interval.microseconds, 1000L));""".stripMargin)
case dt: DecimalType =>
// DecimalOutputShortFastPath: precision <= 18 fits in a signed long, so pass the unscaled
// value to `setSafe(int, long)` and skip the BigDecimal allocation.
Expand Down Expand Up @@ -398,6 +407,7 @@ private[codegen] object CometBatchKernelCodegenOutput {
case ShortType => s"$target.getShort($idx)"
case IntegerType | DateType => s"$target.getInt($idx)"
case LongType | TimestampType | TimestampNTZType => s"$target.getLong($idx)"
case CalendarIntervalType => s"$target.getInterval($idx)"
case FloatType => s"$target.getFloat($idx)"
case DoubleType => s"$target.getDouble($idx)"
case dt: DecimalType => s"$target.getDecimal($idx, ${dt.precision}, ${dt.scale})"
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -43,6 +43,7 @@ private[codegen] object CometSpecializedGettersDispatch {
case IntegerType | DateType => java.lang.Integer.valueOf(g.getInt(ordinal))
case LongType | TimestampType | TimestampNTZType =>
java.lang.Long.valueOf(g.getLong(ordinal))
case CalendarIntervalType => g.getInterval(ordinal)
case FloatType => java.lang.Float.valueOf(g.getFloat(ordinal))
case DoubleType => java.lang.Double.valueOf(g.getDouble(ordinal))
case _: StringType => g.getUTF8String(ordinal)
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -504,7 +504,7 @@ object QueryPlanSerde extends Logging with CometExprShim with CometTypeShim {
def supportedDataType(dt: DataType, allowComplex: Boolean = false): Boolean = dt match {
case _: ByteType | _: ShortType | _: IntegerType | _: LongType | _: FloatType |
_: DoubleType | _: StringType | _: BinaryType | _: TimestampType | _: TimestampNTZType |
_: DecimalType | _: DateType | _: BooleanType | _: NullType =>
_: DecimalType | _: DateType | _: BooleanType | _: NullType | CalendarIntervalType =>
true
case dt if isTimeType(dt) =>
true
Expand Down Expand Up @@ -545,6 +545,7 @@ object QueryPlanSerde extends Logging with CometExprShim with CometTypeShim {
case dt if isTimeType(dt) => 17
case _: YearMonthIntervalType => 18
case _: DayTimeIntervalType => 19
case CalendarIntervalType => 20
case dt =>
logWarning(s"Cannot serialize Spark data type: $dt")
return None
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -83,8 +83,8 @@ private[arrow] object ArrowWriter {
case (_: YearMonthIntervalType, vector: IntervalYearVector) =>
new IntervalYearWriter(vector)
case (_: DayTimeIntervalType, vector: DurationVector) => new DurationWriter(vector)
// case (CalendarIntervalType, vector: IntervalMonthDayNanoVector) =>
// new IntervalMonthDayNanoWriter(vector)
case (CalendarIntervalType, vector: IntervalMonthDayNanoVector) =>
new IntervalMonthDayNanoWriter(vector)
case (dt, _) =>
throw QueryExecutionErrors.notSupportTypeError(dt)
}
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -114,6 +114,8 @@ object Utils extends CometTypeShim with Logging {
case yi: ArrowType.Interval if yi.getUnit == IntervalUnit.YEAR_MONTH =>
YearMonthIntervalType()
case di: ArrowType.Interval if di.getUnit == IntervalUnit.DAY_TIME => DayTimeIntervalType()
case ci: ArrowType.Interval if ci.getUnit == IntervalUnit.MONTH_DAY_NANO =>
CalendarIntervalType
case d: ArrowType.Duration if d.getUnit == TimeUnit.MICROSECOND => DayTimeIntervalType()
case t: ArrowType.Time if t.getUnit == TimeUnit.NANOSECOND && t.getBitWidth == 64 =>
// scalastyle:off classforname
Expand Down Expand Up @@ -162,6 +164,7 @@ object Utils extends CometTypeShim with Logging {
// Spark stores DayTimeIntervalType as microseconds in an int64, matching Arrow
// Duration(Microsecond) rather than the lossy Interval(DayTime) {days, millis} layout.
case _: DayTimeIntervalType => new ArrowType.Duration(TimeUnit.MICROSECOND)
case CalendarIntervalType => new ArrowType.Interval(IntervalUnit.MONTH_DAY_NANO)
case _ =>
throw new UnsupportedOperationException(
s"Unsupported data type: [${dt.getClass.getName}] ${dt.catalogString}")
Expand Down
Original file line number Diff line number Diff line change
@@ -0,0 +1,26 @@
-- 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.

-- Config: spark.comet.exec.localTableScan.enabled=true

query
SELECT * FROM VALUES
(make_interval(1, 2, 3, 4, 5, 6, 7.008009)),
(make_interval(30, 25, 0, -100, 40, 80, 299.889987299)),
(make_interval(0, -1, 0, 1, 0, 0, -1)),
(CAST(NULL AS INTERVAL))
AS test_calendar_interval(i)
Loading
Loading