Skip to content

Commit

Permalink
[#3903] fix(common): convert literal values of partition in `DTOConve…
Browse files Browse the repository at this point in the history
…rters` (#3904)

### What changes were proposed in this pull request?

Convert literal values of partition in `DTOConverters`.

### Why are the changes needed?

Fix: #3903 

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

No.

### How was this patch tested?

UT.

Co-authored-by: zhanghan18 <zhanghan18@xiaomi.com>
  • Loading branch information
xiaozcy and zhanghan18 committed Jun 19, 2024
1 parent 282fc0d commit 5a452f5
Show file tree
Hide file tree
Showing 2 changed files with 33 additions and 12 deletions.
Original file line number Diff line number Diff line change
Expand Up @@ -705,22 +705,33 @@ public static Partition fromDTO(PartitionDTO partitionDTO) {
switch (partitionDTO.type()) {
case IDENTITY:
IdentityPartitionDTO identityPartitionDTO = (IdentityPartitionDTO) partitionDTO;
Literal<?>[] values =
Arrays.stream(identityPartitionDTO.values())
.map(DTOConverters::fromFunctionArg)
.toArray(Literal<?>[]::new);
return Partitions.identity(
identityPartitionDTO.name(),
identityPartitionDTO.fieldNames(),
identityPartitionDTO.values(),
values,
identityPartitionDTO.properties());
case RANGE:
RangePartitionDTO rangePartitionDTO = (RangePartitionDTO) partitionDTO;
return Partitions.range(
rangePartitionDTO.name(),
rangePartitionDTO.upper(),
rangePartitionDTO.lower(),
(Literal<?>) fromFunctionArg(rangePartitionDTO.upper()),
(Literal<?>) fromFunctionArg(rangePartitionDTO.lower()),
rangePartitionDTO.properties());
case LIST:
ListPartitionDTO listPartitionDTO = (ListPartitionDTO) partitionDTO;
return Partitions.list(
listPartitionDTO.name(), listPartitionDTO.lists(), listPartitionDTO.properties());
Literal<?>[][] lists =
Arrays.stream(listPartitionDTO.lists())
.map(
list ->
Arrays.stream(list)
.map(DTOConverters::fromFunctionArg)
.toArray(Literal<?>[]::new))
.toArray(Literal<?>[][]::new);
return Partitions.list(listPartitionDTO.name(), lists, listPartitionDTO.properties());
default:
throw new IllegalArgumentException("Unsupported partition type: " + partitionDTO.type());
}
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -16,6 +16,7 @@
import com.datastrato.gravitino.rel.expressions.literals.Literals;
import com.datastrato.gravitino.rel.expressions.transforms.Transform;
import com.datastrato.gravitino.rel.expressions.transforms.Transforms;
import com.datastrato.gravitino.rel.partitions.IdentityPartition;
import com.datastrato.gravitino.rel.partitions.ListPartition;
import com.datastrato.gravitino.rel.partitions.Partition;
import com.datastrato.gravitino.rel.partitions.Partitions;
Expand All @@ -41,6 +42,10 @@ void testIdentityPartitionDTOConvert() {
LiteralDTO.builder().withDataType(Types.StringType.get()).withValue("us").build();
String[][] fieldNames = {field1, field2};
LiteralDTO[] values = {literal1, literal2};
Literal<?>[] expectedValues = {
(Literal<?>) DTOConverters.fromFunctionArg(literal1),
(Literal<?>) DTOConverters.fromFunctionArg(literal2)
};

Map<String, String> properties = Collections.singletonMap("key", "value");
PartitionDTO identityPartitionDTO =
Expand All @@ -51,14 +56,13 @@ void testIdentityPartitionDTOConvert() {
.withProperties(properties)
.build();
// when
com.datastrato.gravitino.rel.partitions.IdentityPartition identityPartition =
(com.datastrato.gravitino.rel.partitions.IdentityPartition)
DTOConverters.fromDTO(identityPartitionDTO);
IdentityPartition identityPartition =
(IdentityPartition) DTOConverters.fromDTO(identityPartitionDTO);

// then
Assertions.assertTrue(Arrays.equals(fieldNames, identityPartition.fieldNames()));
Assertions.assertEquals("IdentityPartition", identityPartition.name());
Assertions.assertTrue(Arrays.equals(values, identityPartition.values()));
Assertions.assertTrue(Arrays.equals(expectedValues, identityPartition.values()));
Assertions.assertEquals(properties, identityPartition.properties());
}

Expand All @@ -70,6 +74,8 @@ void testRangePartitionDTOConvert() {
LiteralDTO.builder().withDataType(Types.DateType.get()).withValue("2008-08-08").build();
LiteralDTO upper =
LiteralDTO.builder().withDataType(Types.StringType.get()).withValue("us").build();
Literal<?> expectedLower = (Literal<?>) DTOConverters.fromFunctionArg(lower);
Literal<?> expectedUpper = (Literal<?>) DTOConverters.fromFunctionArg(upper);

Map<String, String> properties = Collections.singletonMap("key", "value");
PartitionDTO rangePartitionDTO =
Expand All @@ -84,8 +90,8 @@ void testRangePartitionDTOConvert() {

// then
Assertions.assertEquals("RangePartition", rangePartition.name());
Assertions.assertEquals(lower, rangePartition.lower());
Assertions.assertEquals(upper, rangePartition.upper());
Assertions.assertEquals(expectedLower, rangePartition.lower());
Assertions.assertEquals(expectedUpper, rangePartition.upper());
Assertions.assertEquals(properties, rangePartition.properties());
}

Expand All @@ -100,6 +106,10 @@ void testListPartitionDTOConvert() {

Map<String, String> properties = Collections.singletonMap("key", "value");
LiteralDTO[][] literalDTOS = {new LiteralDTO[] {literal1}, new LiteralDTO[] {literal2}};
Literal<?>[][] expectedValues = {
new Literal<?>[] {(Literal<?>) DTOConverters.fromFunctionArg(literal1)},
new Literal<?>[] {(Literal<?>) DTOConverters.fromFunctionArg(literal2)}
};
ListPartitionDTO listPartitionDTO =
ListPartitionDTO.builder()
.withName("ListPartition")
Expand All @@ -112,7 +122,7 @@ void testListPartitionDTOConvert() {

// then
Assertions.assertEquals("ListPartition", listPartition.name());
Assertions.assertTrue(Arrays.equals(literalDTOS, listPartition.lists()));
Assertions.assertTrue(Arrays.deepEquals(expectedValues, listPartition.lists()));
Assertions.assertEquals(properties, listPartition.properties());
}

Expand Down

0 comments on commit 5a452f5

Please sign in to comment.