Skip to content

Commit

Permalink
[#303]feat(common): Implement Partition JSON SerDe (#304)
Browse files Browse the repository at this point in the history
### What changes were proposed in this pull request?

add partitionDTO

### Why are the changes needed?

Fix: #303 

### Does this PR introduce _any_ user-facing change?
no

### How was this patch tested?

UT added.
  • Loading branch information
mchades authored Sep 4, 2023
1 parent c7777b9 commit 42d859f
Show file tree
Hide file tree
Showing 8 changed files with 730 additions and 51 deletions.
Original file line number Diff line number Diff line change
@@ -0,0 +1,193 @@
/*
* Copyright 2023 Datastrato.
* This software is licensed under the Apache License version 2.
*/
package com.datastrato.graviton.dto.rel;

import static com.datastrato.graviton.dto.rel.ExpressionPartitionDTO.ExpressionType.FIELD;
import static com.datastrato.graviton.dto.rel.ExpressionPartitionDTO.ExpressionType.FUNCTION;
import static com.datastrato.graviton.dto.rel.ExpressionPartitionDTO.ExpressionType.LITERAL;

import com.datastrato.graviton.json.JsonUtils;
import com.fasterxml.jackson.annotation.JsonCreator;
import com.fasterxml.jackson.annotation.JsonProperty;
import com.fasterxml.jackson.annotation.JsonSubTypes;
import com.fasterxml.jackson.annotation.JsonTypeInfo;
import com.fasterxml.jackson.databind.annotation.JsonDeserialize;
import com.fasterxml.jackson.databind.annotation.JsonSerialize;
import com.google.common.base.Preconditions;
import io.substrait.type.Type;
import lombok.EqualsAndHashCode;
import org.apache.logging.log4j.util.Strings;

@EqualsAndHashCode(callSuper = false)
public class ExpressionPartitionDTO implements Partition {

@JsonProperty("expression")
private final Expression expression;

@Override
public Strategy strategy() {
return Strategy.EXPRESSION;
}

@JsonCreator
private ExpressionPartitionDTO(
@JsonProperty("strategy") String strategy,
@JsonProperty("expression") Expression expression) {
Preconditions.checkArgument(expression != null, "expression cannot be null");
this.expression = expression;
}

public static class Builder {
private Expression expression;

public Builder(Expression expression) {
this.expression = expression;
}

public ExpressionPartitionDTO build() {
return new ExpressionPartitionDTO(Strategy.EXPRESSION.name(), expression);
}
}

enum ExpressionType {
FIELD,
LITERAL,
FUNCTION,
}

@JsonTypeInfo(use = JsonTypeInfo.Id.DEDUCTION)
@JsonSubTypes({
@JsonSubTypes.Type(value = ExpressionPartitionDTO.FieldExpression.class),
@JsonSubTypes.Type(value = ExpressionPartitionDTO.LiteralExpression.class),
@JsonSubTypes.Type(value = ExpressionPartitionDTO.FunctionExpression.class),
})
public interface Expression {
@JsonProperty("expressionType")
ExpressionType expressionType();
}

@EqualsAndHashCode
public static class FieldExpression implements Expression {

@JsonProperty("fieldName")
private final String[] fieldName;

@JsonCreator
private FieldExpression(
@JsonProperty("expressionType") String expressionType,
@JsonProperty("fieldName") String[] fieldName) {
Preconditions.checkArgument(
fieldName != null && fieldName.length != 0, "fieldName cannot be null or empty");
this.fieldName = fieldName;
}

@Override
public ExpressionType expressionType() {
return FIELD;
}

public static class Builder {
private String[] fieldName;

public Builder withFieldName(String[] fieldName) {
this.fieldName = fieldName;
return this;
}

public FieldExpression build() {
return new FieldExpression(FIELD.name(), fieldName);
}
}
}

@EqualsAndHashCode
public static class LiteralExpression implements Expression {

@JsonProperty("type")
@JsonSerialize(using = JsonUtils.TypeSerializer.class)
@JsonDeserialize(using = JsonUtils.TypeDeserializer.class)
private final Type type;

@JsonProperty("value")
private final String value;

@JsonCreator
private LiteralExpression(
@JsonProperty("expressionType") String expressionType,
@JsonProperty("type") Type type,
@JsonProperty("value") String value) {
this.type = type;
this.value = value;
}

@Override
public ExpressionType expressionType() {
return LITERAL;
}

public static class Builder {
private Type type;
private String value;

public Builder withType(Type type) {
this.type = type;
return this;
}

public Builder withValue(String value) {
this.value = value;
return this;
}

public LiteralExpression build() {
return new LiteralExpression(LITERAL.name(), type, value);
}
}
}

@EqualsAndHashCode
public static class FunctionExpression implements Expression {

@JsonProperty("funcName")
private final String funcName;

@JsonProperty("args")
private final Expression[] args;

@JsonCreator
private FunctionExpression(
@JsonProperty("expressionType") String expressionType,
@JsonProperty("funcName") String funcName,
@JsonProperty("args") Expression[] args) {
Preconditions.checkArgument(Strings.isNotBlank(funcName), "funcName cannot be null or empty");
this.funcName = funcName;
this.args = args;
}

@Override
public ExpressionType expressionType() {
return FUNCTION;
}

public static class Builder {
private String funcName;
private Expression[] args;

public Builder withFuncName(String funcName) {
this.funcName = funcName;
return this;
}

public Builder withArgs(Expression[] args) {
this.args = args;
return this;
}

public FunctionExpression build() {
return new FunctionExpression(FUNCTION.name(), funcName, args);
}
}
}
}
Original file line number Diff line number Diff line change
@@ -0,0 +1,95 @@
/*
* Copyright 2023 Datastrato.
* This software is licensed under the Apache License version 2.
*/
package com.datastrato.graviton.dto.rel;

import com.fasterxml.jackson.annotation.JsonCreator;
import com.fasterxml.jackson.annotation.JsonProperty;
import com.google.common.base.Preconditions;
import com.google.common.collect.Lists;
import java.util.Arrays;
import java.util.List;
import lombok.EqualsAndHashCode;
import org.apache.logging.log4j.util.Strings;

@EqualsAndHashCode(callSuper = false)
public class ListPartitionDTO implements Partition {
@JsonProperty("fieldNames")
private final String[][] fieldNames;

@JsonProperty("assignments")
private final Assignment[] assignments;

@JsonCreator
private ListPartitionDTO(
@JsonProperty("strategy") String strategy,
@JsonProperty("fieldNames") String[][] fieldNames,
@JsonProperty("assignments") Assignment[] assignments) {
Preconditions.checkArgument(
fieldNames != null && fieldNames.length != 0, "fieldNames cannot be null or empty");

if (assignments != null && assignments.length != 0) {
Preconditions.checkArgument(
Arrays.stream(assignments)
.allMatch(
assignment ->
Arrays.stream(assignment.values)
.allMatch(v -> v.length == fieldNames.length)),
"Assignment values length must be equal to field number");
}

this.fieldNames = fieldNames;
this.assignments = assignments;
}

@Override
public Strategy strategy() {
return Strategy.LIST;
}

@EqualsAndHashCode
public static class Assignment {

@JsonProperty("name")
private final String name;

@JsonProperty("values")
private final String[][] values;

@JsonCreator
private Assignment(
@JsonProperty("name") String name, @JsonProperty("values") String[][] values) {
Preconditions.checkArgument(
!Strings.isBlank(name), "Assignment name cannot be null or empty");
Preconditions.checkArgument(
values != null && values.length != 0, "values cannot be null or empty");
this.name = name;
this.values = values;
}
}

public static class Builder {

private String[][] fieldNames;

private final List<Assignment> assignments = Lists.newArrayList();

public Builder() {}

public Builder withFieldNames(String[][] fieldNames) {
this.fieldNames = fieldNames;
return this;
}

public Builder withAssignment(String partitionName, String[][] values) {
assignments.add(new Assignment(partitionName, values));
return this;
}

public ListPartitionDTO build() {
return new ListPartitionDTO(
Strategy.LIST.name(), fieldNames, assignments.toArray(new Assignment[0]));
}
}
}
Original file line number Diff line number Diff line change
@@ -0,0 +1,40 @@
/*
* Copyright 2023 Datastrato.
* This software is licensed under the Apache License version 2.
*/
package com.datastrato.graviton.dto.rel;

import com.fasterxml.jackson.annotation.JsonProperty;
import com.fasterxml.jackson.annotation.JsonSubTypes;
import com.fasterxml.jackson.annotation.JsonTypeInfo;

@JsonTypeInfo(
use = JsonTypeInfo.Id.NAME,
property = "strategy",
include = JsonTypeInfo.As.EXISTING_PROPERTY,
visible = true)
@JsonSubTypes({
@JsonSubTypes.Type(
value = SimplePartitionDTO.class,
names = {"identity", "year", "month", "day", "hour"}),
@JsonSubTypes.Type(value = ListPartitionDTO.class, name = "list"),
@JsonSubTypes.Type(value = RangePartitionDTO.class, name = "range"),
@JsonSubTypes.Type(value = ExpressionPartitionDTO.class, name = "expression"),
})
public interface Partition {

/** @return The strategy of partitioning */
@JsonProperty("strategy")
Strategy strategy();

enum Strategy {
IDENTITY,
YEAR,
MONTH,
DAY,
HOUR,
LIST,
RANGE,
EXPRESSION,
}
}
Loading

0 comments on commit 42d859f

Please sign in to comment.