Skip to content

Commit

Permalink
add china typeofday function, and unit test
Browse files Browse the repository at this point in the history
  • Loading branch information
@aaronshan
aaronshan committed Jul 20, 2016
1 parent ca46bd2 commit dc076bc
Show file tree
Hide file tree
Showing 4 changed files with 282 additions and 0 deletions.
88 changes: 88 additions & 0 deletions src/main/java/cc/shanruifeng/functions/udfs/scalar/date/ChinaTypeOfDayFunction.java
View file
Original file line number Diff line number Diff line change
@@ -0,0 +1,88 @@
package cc.shanruifeng.functions.udfs.scalar.date;

import cc.shanruifeng.functions.udfs.utils.ConfigUtils;
import com.facebook.presto.operator.Description;
import com.facebook.presto.operator.scalar.ScalarFunction;
import com.facebook.presto.spi.type.StandardTypes;
import com.facebook.presto.type.SqlType;
import io.airlift.slice.Slice;
import io.airlift.slice.Slices;
import org.joda.time.LocalDate;
import org.joda.time.format.DateTimeFormat;
import org.joda.time.format.DateTimeFormatter;

import java.util.Calendar;
import java.util.Map;

import static java.util.concurrent.TimeUnit.DAYS;

/**
* 1: 法定节假日, 2: 正常周末, 3: 正常工作日 4:攒假的工作日
*
* @author ruifeng.shan
* @date 2016-07-15
* @time 14:44
*/
public class ChinaTypeOfDayFunction {
public final static DateTimeFormatter DEFAULT_DATE_FORMATTER = DateTimeFormat.forPattern("yyyy-MM-dd");
public final static Map<String, String> dayMap = ConfigUtils.getDayMap();

private enum DayType {
HOLIDAY("holiday"), WORKDAY("workday");

private String code;
private DayType(String code) {
this.code = code;
}

public String getCode() {
return this.code;
}
}

@ScalarFunction("typeOfDay")
@Description("Returns the type of day from a date string(yyyy-MM-dd)")
@SqlType(StandardTypes.INTEGER)
public static long typeOfDay(@SqlType(StandardTypes.VARCHAR) Slice string) {
if (string == null) {
return -1;
}

String dateStr = string.toStringUtf8();
try {
String value = dayMap.get(dateStr);
if (DayType.HOLIDAY.getCode().equalsIgnoreCase(value)) {
return 1;
} else if (DayType.WORKDAY.getCode().equalsIgnoreCase(value)) {
return 4;
} else {
LocalDate date = LocalDate.parse(string.toStringUtf8(), DEFAULT_DATE_FORMATTER);
if (date.getDayOfWeek() < 6) {
return 3;
} else {
return 2;
}
}
} catch (Exception e) {
}

return -1;
}

@ScalarFunction("typeOfDay")
@Description("Returns the day of week from a date string")
@SqlType(StandardTypes.INTEGER)
public static long typeOfDay(@SqlType(StandardTypes.DATE) long date) {
try {
Calendar calendar = Calendar.getInstance();
calendar.setTimeInMillis(DAYS.toMillis(date));
LocalDate localDate = LocalDate.fromCalendarFields(calendar);

String dateStr = localDate.toString(DEFAULT_DATE_FORMATTER);
return typeOfDay(Slices.utf8Slice(dateStr));
} catch (Exception e) {
}

return -1;
}
}
17 changes: 17 additions & 0 deletions src/main/java/cc/shanruifeng/functions/udfs/utils/ConfigUtils.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -64,4 +64,21 @@ public static Map<String, ChinaIdArea> getIdCardMap() {

return map;
}

public static Map<String, String> getDayMap() {
String fileName = "/china_day_type.config";
Map<String, String> map = Maps.newHashMap();
try {
List<String> list = loadFile(fileName);
for (String line : list) {
String[] results = line.split("\t", 2);
map.put(results[0], results[1]);
}
} catch (IOException e) {
logger.error("get day map error. error is {}.", e);
return map;
}

return map;
}
}
149 changes: 149 additions & 0 deletions src/main/resources/china_day_type.config
View file
Original file line number Diff line number Diff line change
@@ -0,0 +1,149 @@
2013-01-01 holiday
2013-01-02 holiday
2013-01-03 holiday
2013-01-05 workday
2013-01-06 workday
2013-02-09 holiday
2013-02-10 holiday
2013-02-11 holiday
2013-02-12 holiday
2013-02-13 holiday
2013-02-14 holiday
2013-02-15 holiday
2013-02-16 workday
2013-02-17 workday
2013-04-04 holiday
2013-04-05 holiday
2013-04-06 holiday
2013-04-07 workday
2013-04-27 workday
2013-04-28 workday
2013-04-29 holiday
2013-04-30 holiday
2013-05-01 holiday
2013-06-08 workday
2013-06-09 workday
2013-06-10 holiday
2013-06-11 holiday
2013-06-12 holiday
2013-09-19 holiday
2013-09-20 holiday
2013-09-21 holiday
2013-09-22 workday
2013-09-29 workday
2013-10-01 holiday
2013-10-02 holiday
2013-10-03 holiday
2013-10-04 holiday
2013-10-05 holiday
2013-10-06 holiday
2013-10-07 holiday
2013-10-12 workday
2014-01-01 holiday
2014-01-21 holiday
2014-01-26 workday
2014-02-01 holiday
2014-02-02 holiday
2014-02-03 holiday
2014-02-04 holiday
2014-02-05 holiday
2014-02-06 holiday
2014-02-08 workday
2014-04-05 holiday
2014-04-06 holiday
2014-04-07 holiday
2014-05-01 holiday
2014-05-02 holiday
2014-05-03 holiday
2014-05-04 workday
2014-05-31 holiday
2014-06-01 holiday
2014-06-02 holiday
2014-09-06 holiday
2014-09-07 holiday
2014-09-08 holiday
2014-09-28 workday
2014-10-01 holiday
2014-10-02 holiday
2014-10-03 holiday
2014-10-04 holiday
2014-10-05 holiday
2014-10-06 holiday
2014-10-07 holiday
2014-10-11 workday
2015-01-01 holiday
2015-01-02 holiday
2015-01-03 holiday
2015-01-04 workday
2015-02-15 workday
2015-02-18 holiday
2015-02-19 holiday
2015-02-20 holiday
2015-02-21 holiday
2015-02-22 holiday
2015-02-23 holiday
2015-02-24 holiday
2015-02-28 workday
2015-04-04 holiday
2015-04-05 holiday
2015-04-06 holiday
2015-05-01 holiday
2015-05-02 holiday
2015-05-03 holiday
2015-06-20 holiday
2015-06-21 holiday
2015-06-22 holiday
2015-09-03 holiday
2015-09-04 holiday
2015-09-05 holiday
2015-09-06 workday
2015-09-26 holiday
2015-09-27 holiday
2015-10-01 holiday
2015-10-02 holiday
2015-10-03 holiday
2015-10-04 holiday
2015-10-05 holiday
2015-10-06 holiday
2015-10-07 holiday
2015-10-10 workday
2016-01-01 holiday
2016-01-02 holiday
2016-01-03 holiday
2016-01-04 workday
2016-02-06 workday
2016-02-07 holiday
2016-02-08 holiday
2016-02-08 holiday
2016-02-09 holiday
2016-02-10 holiday
2016-02-11 holiday
2016-02-12 holiday
2016-02-13 holiday
2016-02-14 workday
2016-04-02 holiday
2016-04-03 holiday
2016-04-04 holiday
2016-04-30 holiday
2016-05-01 holiday
2016-05-02 holiday
2016-06-09 holiday
2016-06-10 holiday
2016-06-11 holiday
2016-06-12 workday
2016-09-15 holiday
2016-09-16 holiday
2016-09-17 holiday
2016-09-18 workday
2016-10-01 holiday
2016-10-02 holiday
2016-10-03 holiday
2016-10-04 holiday
2016-10-05 holiday
2016-10-06 holiday
2016-10-07 holiday
2016-10-08 workday
2016-10-09 workday
2017-01-01 holiday
2017-01-02 holiday
2017-01-03 holiday
28 changes: 28 additions & 0 deletions src/test/java/cc/shanruifeng/functions/udfs/scalar/date/ChinaTypeOfDayFunctionTest.java
View file
Original file line number Diff line number Diff line change
@@ -0,0 +1,28 @@
package cc.shanruifeng.functions.udfs.scalar.date;

import com.facebook.presto.metadata.FunctionListBuilder;
import com.facebook.presto.type.TypeRegistry;
import io.airlift.slice.Slices;
import junit.framework.Assert;
import org.junit.Test;

/**
* @author ruifeng.shan
* @date 2016-07-15
* @time 14:48
*/
public class ChinaTypeOfDayFunctionTest {
@Test
public void testFunctionCreate() throws Exception {
TypeRegistry typeRegistry = new TypeRegistry();
FunctionListBuilder builder = new FunctionListBuilder(typeRegistry);
builder.scalar(ChinaTypeOfDayFunction.class);
}
@Test
public void testTypeOfDay() {
Assert.assertEquals(1,ChinaTypeOfDayFunction.typeOfDay(Slices.utf8Slice("2016-10-01")));
Assert.assertEquals(2,ChinaTypeOfDayFunction.typeOfDay(Slices.utf8Slice("2016-07-16")));
Assert.assertEquals(3,ChinaTypeOfDayFunction.typeOfDay(Slices.utf8Slice("2016-07-15")));
Assert.assertEquals(4,ChinaTypeOfDayFunction.typeOfDay(Slices.utf8Slice("2016-09-18")));
}
}

0 comments on commit dc076bc

Please sign in to comment.