commit 1d7591b81b9590503ba0598e7c57d68c05dd213b Author: Bharath Krishna Date: Thu May 10 14:43:14 2018 -0700 HIVE-19370 : Retain time in the output for timestamp inputs in add_months UDF diff --git common/src/java/org/apache/hive/common/util/DateUtils.java common/src/java/org/apache/hive/common/util/DateUtils.java index 65f3b9401916abdfa52fbf75d115ba6b61758fb0..6ac91411f8b868abfdcaea8619b2732a57ce0451 100644 --- common/src/java/org/apache/hive/common/util/DateUtils.java +++ common/src/java/org/apache/hive/common/util/DateUtils.java @@ -36,10 +36,21 @@ protected SimpleDateFormat initialValue() { } }; + private static final ThreadLocal DATE_TIME_FORMAT_LOCAL = + ThreadLocal.withInitial(() -> { + SimpleDateFormat simpleDateTimeFormat = new SimpleDateFormat("yyyy-MM-dd HH:mm:ss"); + simpleDateTimeFormat.setLenient(false); + return simpleDateTimeFormat; + }); + public static SimpleDateFormat getDateFormat() { return dateFormatLocal.get(); } + public static SimpleDateFormat getDateTimeFormat() { + return DATE_TIME_FORMAT_LOCAL.get(); + } + public static final int NANOS_PER_SEC = 1000000000; public static final BigDecimal MAX_INT_BD = new BigDecimal(Integer.MAX_VALUE); public static final BigDecimal NANOS_PER_SEC_BD = new BigDecimal(NANOS_PER_SEC); diff --git ql/src/java/org/apache/hadoop/hive/ql/udf/generic/GenericUDFAddMonths.java ql/src/java/org/apache/hadoop/hive/ql/udf/generic/GenericUDFAddMonths.java index dae4b97b4a17e98122431e5fda655fd9f873fdb5..a7365810e99e2d238d23bef50a530f6ef18554a5 100644 --- ql/src/java/org/apache/hadoop/hive/ql/udf/generic/GenericUDFAddMonths.java +++ ql/src/java/org/apache/hadoop/hive/ql/udf/generic/GenericUDFAddMonths.java @@ -22,6 +22,7 @@ import static org.apache.hadoop.hive.serde2.objectinspector.primitive.PrimitiveObjectInspectorUtils.PrimitiveGrouping.STRING_GROUP; import static org.apache.hadoop.hive.serde2.objectinspector.primitive.PrimitiveObjectInspectorUtils.PrimitiveGrouping.VOID_GROUP; +import java.text.SimpleDateFormat; import java.util.Calendar; import java.util.Date; @@ -31,6 +32,7 @@ import org.apache.hadoop.hive.serde2.objectinspector.ConstantObjectInspector; import org.apache.hadoop.hive.serde2.objectinspector.ObjectInspector; import org.apache.hadoop.hive.serde2.objectinspector.ObjectInspectorConverters.Converter; +import org.apache.hadoop.hive.serde2.objectinspector.PrimitiveObjectInspector; import org.apache.hadoop.hive.serde2.objectinspector.PrimitiveObjectInspector.PrimitiveCategory; import org.apache.hadoop.hive.serde2.objectinspector.primitive.PrimitiveObjectInspectorFactory; import org.apache.hadoop.io.Text; @@ -57,7 +59,7 @@ private final Text output = new Text(); private transient Integer numMonthsConst; private transient boolean isNumMonthsConst; - + private transient PrimitiveCategory inputType1; @Override public ObjectInspector initialize(ObjectInspector[] arguments) throws UDFArgumentException { checkArgsSize(arguments, 2, 2); @@ -68,7 +70,17 @@ public ObjectInspector initialize(ObjectInspector[] arguments) throws UDFArgumen checkArgGroups(arguments, 0, inputTypes, STRING_GROUP, DATE_GROUP, VOID_GROUP); checkArgGroups(arguments, 1, inputTypes, NUMERIC_GROUP, VOID_GROUP); - obtainDateConverter(arguments, 0, inputTypes, converters); + PrimitiveObjectInspector inOi = (PrimitiveObjectInspector) arguments[0]; + inputType1 = inOi.getPrimitiveCategory(); + switch (inputType1) { + case TIMESTAMP: + obtainTimestampConverter(arguments, 0, inputTypes, converters); + break; + default: + obtainDateConverter(arguments, 0, inputTypes, converters); + break; + } + obtainIntConverter(arguments, 1, inputTypes, converters); if (arguments[1] instanceof ConstantObjectInspector) { @@ -92,19 +104,41 @@ public Object evaluate(DeferredObject[] arguments) throws HiveException { if (numMonthV == null) { return null; } - + Date date; int numMonthInt = numMonthV.intValue(); - Date date = getDateValue(arguments, 0, inputTypes, converters); + switch (inputType1) { + case TIMESTAMP: + date = new Date(getTimestampValue(arguments, 0, converters).getTime()); + break; + default: + date = getDateValue(arguments, 0, inputTypes, converters); + break; + } + if (date == null) { return null; } addMonth(date, numMonthInt); Date newDate = calendar.getTime(); - output.set(DateUtils.getDateFormat().format(newDate)); + output.set(getDateFormat().format(newDate)); return output; } + private SimpleDateFormat getDateFormat() { + SimpleDateFormat df; + switch (inputType1) { + case TIMESTAMP: + case TIMESTAMPLOCALTZ: + df = DateUtils.getDateTimeFormat(); + break; + default: + df = DateUtils.getDateFormat(); + break; + } + return df; + } + @Override public String getDisplayString(String[] children) { return getStandardDisplayString(getFuncName(), children); diff --git ql/src/test/org/apache/hadoop/hive/ql/udf/generic/TestGenericUDFAddMonths.java ql/src/test/org/apache/hadoop/hive/ql/udf/generic/TestGenericUDFAddMonths.java index af9b6c43c7dafc69c4944eab02894786af306f35..59e823539bf4c5bf551e63207ff46074e6c0c508 100644 --- ql/src/test/org/apache/hadoop/hive/ql/udf/generic/TestGenericUDFAddMonths.java +++ ql/src/test/org/apache/hadoop/hive/ql/udf/generic/TestGenericUDFAddMonths.java @@ -25,10 +25,12 @@ import org.apache.hadoop.hive.ql.udf.generic.GenericUDF.DeferredObject; import org.apache.hadoop.hive.serde2.io.ByteWritable; import org.apache.hadoop.hive.serde2.io.ShortWritable; +import org.apache.hadoop.hive.serde2.io.TimestampWritable; import org.apache.hadoop.hive.serde2.objectinspector.ObjectInspector; import org.apache.hadoop.hive.serde2.objectinspector.primitive.PrimitiveObjectInspectorFactory; import org.apache.hadoop.io.IntWritable; import org.apache.hadoop.io.Text; +import java.sql.Timestamp; public class TestGenericUDFAddMonths extends TestCase { @@ -63,6 +65,17 @@ public void testAddMonthsInt() throws HiveException { runAndVerify("2016-02-29 10:30:00", -1, "2016-01-31", udf); } + public void testAddMonthsTimestamp() throws HiveException { + GenericUDFAddMonths udf = new GenericUDFAddMonths(); + ObjectInspector valueOI0 = PrimitiveObjectInspectorFactory.writableTimestampObjectInspector; + ObjectInspector valueOI1 = PrimitiveObjectInspectorFactory.writableIntObjectInspector; + ObjectInspector[] arguments = {valueOI0, valueOI1}; + + udf.initialize(arguments); + runAndVerify(Timestamp.valueOf("2018-05-10 08:15:12"), 1, "2018-06-10 08:15:12", udf); + runAndVerify(Timestamp.valueOf("2017-12-31 14:15:16"), 2, "2018-02-28 14:15:16", udf); + } + public void testWrongDateStr() throws HiveException { boolean caught = false; try { @@ -150,6 +163,16 @@ private void runAndVerify(String str, int months, String expResult, GenericUDF u assertEquals("add_months() test ", expResult, output != null ? output.toString() : null); } + private void runAndVerify(Timestamp ts, int months, String expResult, GenericUDF udf) + throws HiveException { + DeferredObject valueObj0 = new DeferredJavaObject(new TimestampWritable(ts)); + DeferredObject valueObj1 = new DeferredJavaObject(new IntWritable(months)); + DeferredObject[] args = {valueObj0, valueObj1}; + Text output = (Text) udf.evaluate(args); + assertEquals("add_months() test ", expResult, output != null ? output.toString() : null); + } + + private void runAndVerify(String str, short months, String expResult, GenericUDF udf) throws HiveException { DeferredObject valueObj0 = new DeferredJavaObject(new Text(str));