This is an automated email from the ASF dual-hosted git repository.
jshao pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/gravitino.git
The following commit(s) were added to refs/heads/main by this push:
new 4426f22277 [#5199] feat(client-python): add partitioning DTO serdes
(#8363)
4426f22277 is described below
commit 4426f2227731740ec2c47664919949b1a818e4fc
Author: George T. C. Lai <[email protected]>
AuthorDate: Thu Sep 4 17:40:29 2025 +0800
[#5199] feat(client-python): add partitioning DTO serdes (#8363)
### What changes were proposed in this pull request?
This PR is aimed at implementing the following classes corresponding to
the Java client.
JsonUtils.java
- PartitioningSerializer
- PartitioningDeserializer
### Why are the changes needed?
We need to support table partitioning, bucketing and sort ordering and
indexes
#5199
### Does this PR introduce _any_ user-facing change?
No
### How was this patch tested?
Unit tests
---------
Signed-off-by: George T. C. Lai <[email protected]>
---
.../gravitino/api/types/json_serdes/base.py | 21 +-
.../dto/rel/partitioning/json_serdes/__init__.py | 16 +
.../json_serdes/partitioning_serdes.py | 221 ++++++++
.../gravitino/dto/rel/partitioning/partitioning.py | 5 +-
clients/client-python/gravitino/utils/serdes.py | 11 +
.../unittests/dto/rel/test_partitioning_serdes.py | 580 +++++++++++++++++++++
6 files changed, 845 insertions(+), 9 deletions(-)
diff --git a/clients/client-python/gravitino/api/types/json_serdes/base.py
b/clients/client-python/gravitino/api/types/json_serdes/base.py
index 5dd89a0266..9f44de8462 100644
--- a/clients/client-python/gravitino/api/types/json_serdes/base.py
+++ b/clients/client-python/gravitino/api/types/json_serdes/base.py
@@ -20,23 +20,32 @@ from typing import Generic, TypeVar, Union
from dataclasses_json.core import Json
+from gravitino.api.expressions.distributions.distribution import Distribution
from gravitino.api.expressions.expression import Expression
+from gravitino.api.expressions.indexes.index import Index
+from gravitino.api.expressions.sorts.sort_order import SortOrder
from gravitino.api.types.types import Type
+from gravitino.dto.rel.partitioning.partitioning import Partitioning
from gravitino.dto.rel.partitions.partition_dto import PartitionDTO
-GravitinoTypeT = TypeVar("GravitinoTypeT", bound=Union[Expression, Type,
PartitionDTO])
+_GravitinoTypeT = TypeVar(
+ "_GravitinoTypeT",
+ bound=Union[
+ Expression, Type, Partitioning, PartitionDTO, Distribution, Index,
SortOrder
+ ],
+)
-class JsonSerializable(ABC, Generic[GravitinoTypeT]):
+class JsonSerializable(ABC, Generic[_GravitinoTypeT]):
"""Customized generic Serializer for DataClassJson."""
@classmethod
@abstractmethod
- def serialize(cls, data_type: GravitinoTypeT) -> Json:
+ def serialize(cls, data_type: _GravitinoTypeT) -> Json:
"""To serialize the given `data`.
Args:
- data (GravitinoTypeT): The data to be serialized.
+ data (_GravitinoTypeT): The data to be serialized.
Returns:
Json: The serialized data.
@@ -45,13 +54,13 @@ class JsonSerializable(ABC, Generic[GravitinoTypeT]):
@classmethod
@abstractmethod
- def deserialize(cls, data: Json) -> GravitinoTypeT:
+ def deserialize(cls, data: Json) -> _GravitinoTypeT:
"""To deserialize the given `data`.
Args:
data (Json): The data to be deserialized.
Returns:
- GravitinoTypeT: The deserialized data.
+ _GravitinoTypeT: The deserialized data.
"""
pass
diff --git
a/clients/client-python/gravitino/dto/rel/partitioning/json_serdes/__init__.py
b/clients/client-python/gravitino/dto/rel/partitioning/json_serdes/__init__.py
new file mode 100644
index 0000000000..13a83393a9
--- /dev/null
+++
b/clients/client-python/gravitino/dto/rel/partitioning/json_serdes/__init__.py
@@ -0,0 +1,16 @@
+# 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.
diff --git
a/clients/client-python/gravitino/dto/rel/partitioning/json_serdes/partitioning_serdes.py
b/clients/client-python/gravitino/dto/rel/partitioning/json_serdes/partitioning_serdes.py
new file mode 100644
index 0000000000..2b3a9bbc03
--- /dev/null
+++
b/clients/client-python/gravitino/dto/rel/partitioning/json_serdes/partitioning_serdes.py
@@ -0,0 +1,221 @@
+# 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.
+
+from contextlib import suppress
+from types import MappingProxyType
+from typing import Any, Dict, Final, cast
+
+from gravitino.api.types.json_serdes.base import JsonSerializable
+from gravitino.dto.rel.expressions.json_serdes._helper.serdes_utils import (
+ SerdesUtils as ExpressionSerdesUtils,
+)
+from gravitino.dto.rel.partitioning.bucket_partitioning_dto import (
+ BucketPartitioningDTO,
+)
+from gravitino.dto.rel.partitioning.day_partitioning_dto import
DayPartitioningDTO
+from gravitino.dto.rel.partitioning.function_partitioning_dto import (
+ FunctionPartitioningDTO,
+)
+from gravitino.dto.rel.partitioning.hour_partitioning_dto import
HourPartitioningDTO
+from gravitino.dto.rel.partitioning.identity_partitioning_dto import (
+ IdentityPartitioningDTO,
+)
+from gravitino.dto.rel.partitioning.list_partitioning_dto import
ListPartitioningDTO
+from gravitino.dto.rel.partitioning.month_partitioning_dto import
MonthPartitioningDTO
+from gravitino.dto.rel.partitioning.partitioning import (
+ Partitioning,
+ SingleFieldPartitioning,
+)
+from gravitino.dto.rel.partitioning.range_partitioning_dto import
RangePartitioningDTO
+from gravitino.dto.rel.partitioning.truncate_partitioning_dto import (
+ TruncatePartitioningDTO,
+)
+from gravitino.dto.rel.partitioning.year_partitioning_dto import
YearPartitioningDTO
+from gravitino.dto.rel.partitions.json_serdes._helper.serdes_utils import (
+ SerdesUtils as PartitionSerdesUtils,
+)
+from gravitino.dto.rel.partitions.list_partition_dto import ListPartitionDTO
+from gravitino.dto.rel.partitions.range_partition_dto import RangePartitionDTO
+from gravitino.utils.precondition import Precondition
+from gravitino.utils.serdes import SerdesUtilsBase
+
+
+class PartitioningSerdes(SerdesUtilsBase, JsonSerializable[Partitioning]):
+ """Custom JSON serializer/deserializer for PartitionDTO objects."""
+
+ _SINGLE_FIELD_PARTITIONING: Final[MappingProxyType] = MappingProxyType(
+ {
+ Partitioning.Strategy.IDENTITY: IdentityPartitioningDTO,
+ Partitioning.Strategy.YEAR: YearPartitioningDTO,
+ Partitioning.Strategy.MONTH: MonthPartitioningDTO,
+ Partitioning.Strategy.DAY: DayPartitioningDTO,
+ Partitioning.Strategy.HOUR: HourPartitioningDTO,
+ }
+ )
+
+ @classmethod
+ def serialize(cls, data_type: Partitioning) -> Dict[str, Any]:
+ """Serialize the given PartitionDTO object.
+
+ Args:
+ data_type (Partitioning): The PartitionDTO objects.
+
+ Returns:
+ Dict[str, Any]: The serialized result.
+
+ Raises:
+ IOError: If partitioning strategy is unknown.
+ """
+
+ strategy = data_type.strategy()
+ result = {cls.STRATEGY: strategy.name.lower()}
+
+ if strategy in cls._SINGLE_FIELD_PARTITIONING:
+ dto = cast(SingleFieldPartitioning, data_type)
+ return {**result, cls.FIELD_NAME: dto.field_name()}
+ if strategy is Partitioning.Strategy.BUCKET:
+ dto = cast(BucketPartitioningDTO, data_type)
+ return {
+ **result,
+ cls.NUM_BUCKETS: dto.num_buckets(),
+ cls.FIELD_NAMES: dto.field_names(),
+ }
+ if strategy is Partitioning.Strategy.TRUNCATE:
+ dto = cast(TruncatePartitioningDTO, data_type)
+ return {
+ **result,
+ cls.WIDTH: dto.width(),
+ cls.FIELD_NAME: dto.field_name(),
+ }
+ if strategy is Partitioning.Strategy.LIST:
+ dto = cast(ListPartitioningDTO, data_type)
+ return {
+ **result,
+ cls.FIELD_NAMES: dto.field_names(),
+ cls.ASSIGNMENTS_NAME: [
+ PartitionSerdesUtils.write_partition(list_partition_dto)
+ for list_partition_dto in dto.assignments()
+ ],
+ }
+ if strategy is Partitioning.Strategy.RANGE:
+ dto = cast(RangePartitioningDTO, data_type)
+ return {
+ **result,
+ cls.FIELD_NAME: dto.field_name(),
+ cls.ASSIGNMENTS_NAME: [
+ PartitionSerdesUtils.write_partition(range_partition_dto)
+ for range_partition_dto in dto.assignments()
+ ],
+ }
+ if strategy is Partitioning.Strategy.FUNCTION:
+ dto = cast(FunctionPartitioningDTO, data_type)
+ return {
+ **result,
+ cls.FUNCTION_NAME: dto.function_name(),
+ cls.FUNCTION_ARGS: [
+ ExpressionSerdesUtils.write_function_arg(function_arg)
+ for function_arg in dto.args()
+ ],
+ }
+
+ raise IOError(f"Unknown partitioning strategy: {strategy}")
+
+ @classmethod
+ def deserialize(cls, data: Dict[str, Any]) -> Partitioning:
+ """Deserialize the given JSON data to the corresponding Partitioning
object.
+
+ Args:
+ data (Dict[str, Any]): The JSON data object to be deserialized.
+
+ Returns:
+ Partitioning: The deserialized result.
+
+ Raises:
+ IllegalArgumentException: If there's illegal arguments in the
given JSON data.
+ """
+
+ Precondition.check_argument(
+ isinstance(data, dict) and len(data) > 0,
+ f"Cannot parse partitioning from invalid JSON: {data}",
+ )
+ Precondition.check_argument(
+ cls.STRATEGY in data,
+ f"Cannot parse partitioning from missing strategy: {data}",
+ )
+ strategy = None
+ with suppress(ValueError):
+ strategy = Partitioning.Strategy(data[cls.STRATEGY].lower())
+
+ if strategy in cls._SINGLE_FIELD_PARTITIONING:
+ return cls._SINGLE_FIELD_PARTITIONING[strategy](
+ *data.get(cls.FIELD_NAME, [])
+ )
+ if strategy is Partitioning.Strategy.BUCKET:
+ return BucketPartitioningDTO(
+ int(data[cls.NUM_BUCKETS]),
+ *data.get(cls.FIELD_NAMES, []),
+ )
+ if strategy is Partitioning.Strategy.TRUNCATE:
+ return TruncatePartitioningDTO(
+ int(data[cls.WIDTH]),
+ data.get(cls.FIELD_NAME, []),
+ )
+ if strategy is Partitioning.Strategy.LIST:
+ field_names = data[cls.FIELD_NAMES]
+ assignments_data = data.get(cls.ASSIGNMENTS_NAME, [])
+ Precondition.check_argument(
+ isinstance(assignments_data, list),
+ f"Cannot parse list partitioning from non-array assignments:
{assignments_data}",
+ )
+ assignments = []
+ for assignment in assignments_data:
+ partition_dto = PartitionSerdesUtils.read_partition(assignment)
+ Precondition.check_argument(
+ isinstance(partition_dto, ListPartitionDTO),
+ f"Cannot parse list partitioning from non-list assignment:
{assignment}",
+ )
+ assignments.append(partition_dto)
+ return ListPartitioningDTO(field_names, assignments)
+ if strategy is Partitioning.Strategy.RANGE:
+ fields = data[cls.FIELD_NAME]
+ assignments_data = data.get(cls.ASSIGNMENTS_NAME, [])
+ Precondition.check_argument(
+ isinstance(assignments_data, list),
+ f"Cannot parse range partitioning from non-array assignments:
{assignments_data}",
+ )
+ assignments = []
+ for assignment in assignments_data:
+ partition_dto = PartitionSerdesUtils.read_partition(assignment)
+ Precondition.check_argument(
+ isinstance(partition_dto, RangePartitionDTO),
+ f"Cannot parse range partitioning from non-range
assignment: {assignment}",
+ )
+ assignments.append(partition_dto)
+ return RangePartitioningDTO(fields, assignments)
+ if strategy is Partitioning.Strategy.FUNCTION:
+ function_name = data[cls.FUNCTION_NAME]
+ Precondition.check_argument(
+ data.get(cls.FUNCTION_ARGS) is not None,
+ f"Cannot parse function partitioning from missing function
args: {data}",
+ )
+ args = [
+ ExpressionSerdesUtils.read_function_arg(arg)
+ for arg in data[cls.FUNCTION_ARGS]
+ ]
+ return FunctionPartitioningDTO(function_name, *args)
+
+ raise IOError(f"Unknown partitioning strategy: {data[cls.STRATEGY]}")
diff --git
a/clients/client-python/gravitino/dto/rel/partitioning/partitioning.py
b/clients/client-python/gravitino/dto/rel/partitioning/partitioning.py
index 600090c469..f505cfcdbb 100644
--- a/clients/client-python/gravitino/dto/rel/partitioning/partitioning.py
+++ b/clients/client-python/gravitino/dto/rel/partitioning/partitioning.py
@@ -22,7 +22,6 @@ from typing import Final, List
from gravitino.api.expressions.expression import Expression
from gravitino.api.expressions.named_reference import NamedReference
from gravitino.api.expressions.transforms.transform import Transform
-from gravitino.dto.rel.column_dto import ColumnDTO
from gravitino.dto.rel.partition_utils import PartitionUtils
from gravitino.exceptions.base import IllegalArgumentException
from gravitino.utils.precondition import Precondition
@@ -76,7 +75,7 @@ class Partitioning(Transform):
pass # pragma: no cover
@abstractmethod
- def validate(self, columns: List[ColumnDTO]) -> None:
+ def validate(self, columns: List["ColumnDTO"]) -> None:
"""Validates the partitioning columns.
Args:
@@ -128,7 +127,7 @@ class SingleFieldPartitioning(Partitioning):
"""
return self._field_name
- def validate(self, columns: List[ColumnDTO]) -> None:
+ def validate(self, columns: List["ColumnDTO"]) -> None:
"""Validates the partitioning columns.
Args:
diff --git a/clients/client-python/gravitino/utils/serdes.py
b/clients/client-python/gravitino/utils/serdes.py
index ed3371dbe8..8554e65cc4 100644
--- a/clients/client-python/gravitino/utils/serdes.py
+++ b/clients/client-python/gravitino/utils/serdes.py
@@ -59,6 +59,17 @@ class SerdesUtilsBase:
LIST_PARTITION_LISTS: Final[str] = "lists"
RANGE_PARTITION_UPPER: Final[str] = "upper"
RANGE_PARTITION_LOWER: Final[str] = "lower"
+ STRATEGY: Final[str] = "strategy"
+ NUM_BUCKETS: Final[str] = "numBuckets"
+ WIDTH: Final[str] = "width"
+ ASSIGNMENTS_NAME: Final[str] = "assignments"
+ SORT_TERM: Final[str] = "sortTerm"
+ DIRECTION: Final[str] = "direction"
+ NULL_ORDERING: Final[str] = "nullOrdering"
+ INDEX_TYPE: Final[str] = "indexType"
+ INDEX_NAME: Final[str] = "name"
+ INDEX_FIELD_NAMES: Final[str] = "fieldNames"
+ NUMBER: Final[str] = "number"
NON_PRIMITIVE_TYPES: Final[Set[Name]] = {
Name.STRUCT,
diff --git
a/clients/client-python/tests/unittests/dto/rel/test_partitioning_serdes.py
b/clients/client-python/tests/unittests/dto/rel/test_partitioning_serdes.py
new file mode 100644
index 0000000000..ee99c7112f
--- /dev/null
+++ b/clients/client-python/tests/unittests/dto/rel/test_partitioning_serdes.py
@@ -0,0 +1,580 @@
+# 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.
+
+import json
+import unittest
+from enum import Enum
+from unittest.mock import patch
+
+from gravitino.api.types.types import Types
+from gravitino.dto.rel.expressions.field_reference_dto import FieldReferenceDTO
+from gravitino.dto.rel.expressions.literal_dto import LiteralDTO
+from gravitino.dto.rel.partitioning.day_partitioning_dto import
DayPartitioningDTO
+from gravitino.dto.rel.partitioning.function_partitioning_dto import (
+ FunctionPartitioningDTO,
+)
+from gravitino.dto.rel.partitioning.hour_partitioning_dto import
HourPartitioningDTO
+from gravitino.dto.rel.partitioning.identity_partitioning_dto import (
+ IdentityPartitioningDTO,
+)
+from gravitino.dto.rel.partitioning.json_serdes.partitioning_serdes import (
+ PartitioningSerdes,
+)
+from gravitino.dto.rel.partitioning.month_partitioning_dto import
MonthPartitioningDTO
+from gravitino.dto.rel.partitioning.partitioning import Partitioning
+from gravitino.dto.rel.partitioning.year_partitioning_dto import
YearPartitioningDTO
+from gravitino.exceptions.base import IllegalArgumentException
+
+
+class MockPartitionStrategy(str, Enum):
+ INVALID_STRATEGY = "invalid_partitioning_strategy"
+
+
+class TestPartitioningSerdes(unittest.TestCase):
+ @classmethod
+ def setUpClass(cls):
+ cls.field_name = [f"dummy_field_{i}" for i in range(1)]
+ cls.single_field_partitioning_dtos = {
+ Partitioning.Strategy.IDENTITY:
IdentityPartitioningDTO(*cls.field_name),
+ Partitioning.Strategy.YEAR: YearPartitioningDTO(*cls.field_name),
+ Partitioning.Strategy.MONTH: MonthPartitioningDTO(*cls.field_name),
+ Partitioning.Strategy.DAY: DayPartitioningDTO(*cls.field_name),
+ Partitioning.Strategy.HOUR: HourPartitioningDTO(*cls.field_name),
+ }
+
+ def test_serialize_invalid_strategy(self):
+ mock_dto = IdentityPartitioningDTO(*self.field_name)
+ with patch.object(
+ mock_dto, "strategy",
return_value=MockPartitionStrategy.INVALID_STRATEGY
+ ):
+ self.assertRaisesRegex(
+ IOError,
+ "Unknown partitioning strategy",
+ PartitioningSerdes.serialize,
+ mock_dto,
+ )
+
+ def test_deserialize_invalid_json(self):
+ invalid_partitioning_data = (None, "invalid_data")
+
+ for invalid_data in invalid_partitioning_data:
+ with self.assertRaisesRegex(
+ IllegalArgumentException, "Cannot parse partitioning from
invalid JSON"
+ ):
+ PartitioningSerdes.deserialize(invalid_data)
+
+ invalid_json_string = "{}"
+ with self.assertRaisesRegex(
+ IllegalArgumentException, "Cannot parse partitioning from invalid
JSON"
+ ):
+ PartitioningSerdes.deserialize(json.loads(invalid_json_string))
+
+ def test_deserialize_invalid_strategy(self):
+ """Tests missing strategy and unknown partitioning strategy."""
+
+ missing_strategy_json_string = """
+ {
+ "fieldName": ["dummy_field_0"]
+ }
+ """
+
+ invalid_strategy_json_string = """
+ {
+ "strategy": "invalid_strategy",
+ "fieldName": ["dummy_field_0"]
+ }
+ """
+
+ with self.assertRaisesRegex(
+ IllegalArgumentException,
+ "Cannot parse partitioning from missing strategy",
+ ):
+
PartitioningSerdes.deserialize(json.loads(missing_strategy_json_string))
+
+ with self.assertRaisesRegex(
+ IOError,
+ "Unknown partitioning strategy",
+ ):
+
PartitioningSerdes.deserialize(json.loads(invalid_strategy_json_string))
+
+ def test_serialize_single_field_partitioning_dto(self):
+ for partitioning_dto in self.single_field_partitioning_dtos.values():
+ serialized = PartitioningSerdes.serialize(partitioning_dto)
+ self.assertEqual(
+ partitioning_dto.name(),
serialized[PartitioningSerdes.STRATEGY]
+ )
+ self.assertEqual(
+ partitioning_dto.strategy().value,
+ serialized[PartitioningSerdes.STRATEGY],
+ )
+ self.assertListEqual(
+ partitioning_dto.field_name(),
+ serialized[PartitioningSerdes.FIELD_NAME],
+ )
+
+ def test_deserialize_identity_partitioning_dto(self):
+ partitioning_dto = self.single_field_partitioning_dtos[
+ Partitioning.Strategy.IDENTITY
+ ]
+ json_string = """
+ {
+ "strategy": "identity",
+ "fieldName": ["dummy_field_0"]
+ }
+ """
+ deserialized = PartitioningSerdes.deserialize(json.loads(json_string))
+
+ self.assertEqual(partitioning_dto.name(), deserialized.name())
+ self.assertEqual(
+ partitioning_dto.strategy().value, deserialized.strategy().value
+ )
+ self.assertListEqual(partitioning_dto.field_name(),
deserialized.field_name())
+
+ def test_deserialize_year_partitioning_dto(self):
+ partitioning_dto = self.single_field_partitioning_dtos[
+ Partitioning.Strategy.YEAR
+ ]
+ json_string = """
+ {
+ "strategy": "year",
+ "fieldName": ["dummy_field_0"]
+ }
+ """
+ deserialized = PartitioningSerdes.deserialize(json.loads(json_string))
+
+ self.assertEqual(partitioning_dto.name(), deserialized.name())
+ self.assertEqual(
+ partitioning_dto.strategy().value, deserialized.strategy().value
+ )
+ self.assertListEqual(partitioning_dto.field_name(),
deserialized.field_name())
+
+ def test_deserialize_month_partitioning_dto(self):
+ partitioning_dto = self.single_field_partitioning_dtos[
+ Partitioning.Strategy.MONTH
+ ]
+ json_string = """
+ {
+ "strategy": "month",
+ "fieldName": ["dummy_field_0"]
+ }
+ """
+ deserialized = PartitioningSerdes.deserialize(json.loads(json_string))
+
+ self.assertEqual(partitioning_dto.name(), deserialized.name())
+ self.assertEqual(
+ partitioning_dto.strategy().value, deserialized.strategy().value
+ )
+ self.assertListEqual(partitioning_dto.field_name(),
deserialized.field_name())
+
+ def test_deserialize_day_partitioning_dto(self):
+ partitioning_dto = self.single_field_partitioning_dtos[
+ Partitioning.Strategy.DAY
+ ]
+ json_string = """
+ {
+ "strategy": "day",
+ "fieldName": ["dummy_field_0"]
+ }
+ """
+ deserialized = PartitioningSerdes.deserialize(json.loads(json_string))
+
+ self.assertEqual(partitioning_dto.name(), deserialized.name())
+ self.assertEqual(
+ partitioning_dto.strategy().value, deserialized.strategy().value
+ )
+ self.assertListEqual(partitioning_dto.field_name(),
deserialized.field_name())
+
+ def test_deserialize_hour_partitioning_dto(self):
+ partitioning_dto = self.single_field_partitioning_dtos[
+ Partitioning.Strategy.HOUR
+ ]
+ json_string = """
+ {
+ "strategy": "hour",
+ "fieldName": ["dummy_field_0"]
+ }
+ """
+ deserialized = PartitioningSerdes.deserialize(json.loads(json_string))
+
+ self.assertEqual(partitioning_dto.name(), deserialized.name())
+ self.assertEqual(
+ partitioning_dto.strategy().value, deserialized.strategy().value
+ )
+ self.assertListEqual(partitioning_dto.field_name(),
deserialized.field_name())
+
+ def test_serdes_bucket_partitioning_dto(self):
+ field_names = [["score"]]
+ json_string = """
+ {
+ "strategy": "bucket",
+ "numBuckets": 10,
+ "fieldNames": [["score"]]
+ }
+ """
+
+ expected_serialized = json.loads(json_string)
+ deserialized = PartitioningSerdes.deserialize(expected_serialized)
+
+ self.assertEqual(Partitioning.Strategy.BUCKET.name.lower(),
deserialized.name())
+ self.assertEqual(
+ Partitioning.Strategy.BUCKET.value, deserialized.strategy().value
+ )
+ self.assertListEqual(field_names, deserialized.field_names())
+
+ serialized = PartitioningSerdes.serialize(deserialized)
+ self.assertDictEqual(expected_serialized, serialized)
+
+ def test_serdes_truncate_partitioning_dto(self):
+ field_name = ["score"]
+ json_string = """
+ {
+ "strategy": "truncate",
+ "width": 20,
+ "fieldName": ["score"]
+ }
+ """
+
+ expected_serialized = json.loads(json_string)
+ deserialized = PartitioningSerdes.deserialize(expected_serialized)
+
+ self.assertEqual(
+ Partitioning.Strategy.TRUNCATE.name.lower(), deserialized.name()
+ )
+ self.assertEqual(
+ Partitioning.Strategy.TRUNCATE.value, deserialized.strategy().value
+ )
+ self.assertListEqual(field_name, deserialized.field_name())
+
+ serialized = PartitioningSerdes.serialize(deserialized)
+ self.assertDictEqual(expected_serialized, serialized)
+
+ def test_serdes_list_partitioning_dto_invalid_assignments(self):
+ json_string = """
+ {
+ "strategy": "list",
+ "fieldNames": [["createTime"], ["city"]],
+ "assignments": "invalid_assignments"
+ }
+ """
+ with self.assertRaisesRegex(
+ IllegalArgumentException,
+ "Cannot parse list partitioning from non-array assignments",
+ ):
+ PartitioningSerdes.deserialize(json.loads(json_string))
+
+ def test_serdes_list_partitioning_dto_invalid_list_assignment(self):
+ json_string = """
+ {
+ "strategy": "list",
+ "fieldNames": [["createTime"], ["city"]],
+ "assignments": [
+ {
+ "type": "range",
+ "name": "p20200321",
+ "upper": {
+ "type": "literal",
+ "dataType": "date",
+ "value": "2020-03-21"
+ },
+ "lower": {
+ "type": "literal",
+ "dataType": "null",
+ "value": "null"
+ }
+ }
+ ]
+ }
+ """
+ with self.assertRaisesRegex(
+ IllegalArgumentException,
+ "Cannot parse list partitioning from non-list assignment",
+ ):
+ PartitioningSerdes.deserialize(json.loads(json_string))
+
+ def test_serdes_list_partitioning_dto(self):
+ field_names = [["createTime"], ["city"]]
+ json_string = """
+ {
+ "strategy": "list",
+ "fieldNames": [["createTime"], ["city"]]
+ }
+ """
+
+ expected_serialized = json.loads(json_string)
+ deserialized = PartitioningSerdes.deserialize(expected_serialized)
+
+ self.assertEqual(Partitioning.Strategy.LIST.name.lower(),
deserialized.name())
+ self.assertEqual(
+ Partitioning.Strategy.LIST.value, deserialized.strategy().value
+ )
+ self.assertListEqual(field_names, deserialized.field_names())
+
+ serialized = PartitioningSerdes.serialize(deserialized)
+ self.assertEqual(
+ expected_serialized[PartitioningSerdes.STRATEGY],
+ serialized[PartitioningSerdes.STRATEGY],
+ )
+ self.assertListEqual(
+ expected_serialized[PartitioningSerdes.FIELD_NAMES],
+ serialized[PartitioningSerdes.FIELD_NAMES],
+ )
+ self.assertEqual(
+ [],
+ serialized[PartitioningSerdes.ASSIGNMENTS_NAME],
+ )
+
+ json_string = """
+ {
+ "strategy": "list",
+ "fieldNames": [["createTime"], ["city"]],
+ "assignments": [
+ {
+ "type": "list",
+ "name": "p202204_California",
+ "properties": {},
+ "lists": [
+ [
+ {
+ "type": "literal",
+ "dataType": "date",
+ "value": "2022-04-01"
+ },
+ {
+ "type": "literal",
+ "dataType": "string",
+ "value": "Los Angeles"
+ }
+ ],
+ [
+ {
+ "type": "literal",
+ "dataType": "date",
+ "value": "2022-04-01"
+ },
+ {
+ "type": "literal",
+ "dataType": "string",
+ "value": "San Francisco"
+ }
+ ]
+ ]
+ }
+ ]
+ }
+ """
+
+ expected_serialized = json.loads(json_string)
+ deserialized = PartitioningSerdes.deserialize(expected_serialized)
+ serialized = PartitioningSerdes.serialize(deserialized)
+ self.assertEqual(
+ expected_serialized[PartitioningSerdes.STRATEGY],
+ serialized[PartitioningSerdes.STRATEGY],
+ )
+ self.assertListEqual(
+ expected_serialized[PartitioningSerdes.FIELD_NAMES],
+ serialized[PartitioningSerdes.FIELD_NAMES],
+ )
+ self.assertEqual(
+ expected_serialized[PartitioningSerdes.ASSIGNMENTS_NAME],
+ serialized[PartitioningSerdes.ASSIGNMENTS_NAME],
+ )
+
+ def test_serdes_range_partitioning_dto_invalid_assignments(self):
+ json_string = """
+ {
+ "strategy": "range",
+ "fieldName": ["dummy_field_0"],
+ "assignments": "invalid_assignments"
+ }
+ """
+ with self.assertRaisesRegex(
+ IllegalArgumentException,
+ "Cannot parse range partitioning from non-array assignments",
+ ):
+ PartitioningSerdes.deserialize(json.loads(json_string))
+
+ def test_serdes_range_partitioning_dto_invalid_range_assignment(self):
+ json_string = """
+ {
+ "strategy": "range",
+ "fieldName": ["dummy_field_0"],
+ "assignments": [
+ {
+ "type": "list",
+ "name": "p202204_California",
+ "properties": {},
+ "lists": [
+ [
+ {
+ "type": "literal",
+ "dataType": "date",
+ "value": "2022-04-01"
+ },
+ {
+ "type": "literal",
+ "dataType": "string",
+ "value": "Los Angeles"
+ }
+ ]
+ ]
+ }
+ ]
+ }
+ """
+ with self.assertRaisesRegex(
+ IllegalArgumentException,
+ "Cannot parse range partitioning from non-range assignment",
+ ):
+ PartitioningSerdes.deserialize(json.loads(json_string))
+
+ def test_serdes_range_partitioning_dto(self):
+ json_string = """
+ {
+ "strategy": "range",
+ "fieldName": ["dummy_field_0"]
+ }
+ """
+
+ expected_serialized = json.loads(json_string)
+ deserialized = PartitioningSerdes.deserialize(expected_serialized)
+
+ self.assertEqual(Partitioning.Strategy.RANGE.name.lower(),
deserialized.name())
+ self.assertEqual(
+ Partitioning.Strategy.RANGE.value, deserialized.strategy().value
+ )
+ self.assertListEqual(
+ TestPartitioningSerdes.field_name, deserialized.field_name()
+ )
+
+ serialized = PartitioningSerdes.serialize(deserialized)
+ self.assertEqual(
+ expected_serialized[PartitioningSerdes.STRATEGY],
+ serialized[PartitioningSerdes.STRATEGY],
+ )
+ self.assertListEqual(
+ expected_serialized[PartitioningSerdes.FIELD_NAME],
+ serialized[PartitioningSerdes.FIELD_NAME],
+ )
+ self.assertEqual(
+ [],
+ serialized[PartitioningSerdes.ASSIGNMENTS_NAME],
+ )
+
+ json_string = """
+ {
+ "strategy": "range",
+ "fieldName": ["dummy_field_0"],
+ "assignments": [
+ {
+ "type": "range",
+ "name": "p20200321",
+ "upper": {
+ "type": "literal",
+ "dataType": "date",
+ "value": "2020-03-21"
+ },
+ "lower": {
+ "type": "literal",
+ "dataType": "null",
+ "value": "null"
+ },
+ "properties": {"key": "value"}
+ }
+ ]
+ }
+ """
+
+ expected_serialized = json.loads(json_string)
+ deserialized = PartitioningSerdes.deserialize(expected_serialized)
+ serialized = PartitioningSerdes.serialize(deserialized)
+ self.assertEqual(
+ expected_serialized[PartitioningSerdes.STRATEGY],
+ serialized[PartitioningSerdes.STRATEGY],
+ )
+ self.assertListEqual(
+ expected_serialized[PartitioningSerdes.FIELD_NAME],
+ serialized[PartitioningSerdes.FIELD_NAME],
+ )
+ self.assertEqual(
+ expected_serialized[PartitioningSerdes.ASSIGNMENTS_NAME],
+ serialized[PartitioningSerdes.ASSIGNMENTS_NAME],
+ )
+
+ def test_serdes_function_partitioning_dto_invalid_args(self):
+ json_string = """
+ {
+ "strategy": "function",
+ "funcName": "dummy_func_name"
+ }
+ """
+
+ with self.assertRaisesRegex(
+ IllegalArgumentException,
+ "Cannot parse function partitioning from missing function args",
+ ):
+ PartitioningSerdes.deserialize(json.loads(json_string))
+
+ def test_serdes_function_partitioning_dto(self):
+ field_arg = FieldReferenceDTO.builder().with_column_name("dt").build()
+ literal_arg = (
+ LiteralDTO.builder()
+ .with_data_type(Types.StringType.get())
+ .with_value("Asia/Shanghai")
+ .build()
+ )
+ json_string = """
+ {
+ "strategy": "function",
+ "funcName": "to_date",
+ "funcArgs": [
+ {
+ "type": "field",
+ "fieldName": ["dt"]
+ },
+ {
+ "type": "literal",
+ "dataType": "string",
+ "value": "Asia/Shanghai"
+ }
+ ]
+ }
+ """
+
+ expected_serialized = json.loads(json_string)
+ deserialized = PartitioningSerdes.deserialize(expected_serialized)
+
+ self.assertIsInstance(deserialized, FunctionPartitioningDTO)
+ self.assertEqual("to_date", deserialized.function_name())
+ self.assertEqual(deserialized.function_name(), deserialized.name())
+ self.assertEqual(
+ Partitioning.Strategy.FUNCTION.value, deserialized.strategy().value
+ )
+ self.assertListEqual([field_arg, literal_arg], deserialized.args())
+
+ serialized = PartitioningSerdes.serialize(deserialized)
+ self.assertEqual(
+ expected_serialized[PartitioningSerdes.STRATEGY],
+ serialized[PartitioningSerdes.STRATEGY],
+ )
+ self.assertEqual(
+ expected_serialized[PartitioningSerdes.FUNCTION_NAME],
+ serialized[PartitioningSerdes.FUNCTION_NAME],
+ )
+ self.assertListEqual(
+ expected_serialized[PartitioningSerdes.FUNCTION_ARGS],
+ serialized[PartitioningSerdes.FUNCTION_ARGS],
+ )
