From e54c6d5581b32f505de431b209d728851f7b6177 Mon Sep 17 00:00:00 2001 From: SophieTech88 Date: Tue, 7 Jan 2025 02:54:17 -0800 Subject: [PATCH] [#5730] feat(client-python): Add sorts expression (#5879) ### What changes were proposed in this pull request? Implement sorts expression in python client, add unit test. ### Why are the changes needed? We need to support the sorts expressions in python client Fix: #5730 ### Does this PR introduce _any_ user-facing change? No ### How was this patch tested? Need to pass all unit tests. --------- Co-authored-by: Xun Co-authored-by: Xun --- .../api/expressions/sorts/__init__.py | 16 +++ .../api/expressions/sorts/null_ordering.py | 35 ++++++ .../api/expressions/sorts/sort_direction.py | 73 +++++++++++ .../api/expressions/sorts/sort_order.py | 45 +++++++ .../api/expressions/sorts/sort_orders.py | 94 ++++++++++++++ .../tests/unittests/rel/test_sorts.py | 118 ++++++++++++++++++ 6 files changed, 381 insertions(+) create mode 100644 clients/client-python/gravitino/api/expressions/sorts/__init__.py create mode 100644 clients/client-python/gravitino/api/expressions/sorts/null_ordering.py create mode 100644 clients/client-python/gravitino/api/expressions/sorts/sort_direction.py create mode 100644 clients/client-python/gravitino/api/expressions/sorts/sort_order.py create mode 100644 clients/client-python/gravitino/api/expressions/sorts/sort_orders.py create mode 100644 clients/client-python/tests/unittests/rel/test_sorts.py diff --git a/clients/client-python/gravitino/api/expressions/sorts/__init__.py b/clients/client-python/gravitino/api/expressions/sorts/__init__.py new file mode 100644 index 00000000000..13a83393a91 --- /dev/null +++ b/clients/client-python/gravitino/api/expressions/sorts/__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/api/expressions/sorts/null_ordering.py b/clients/client-python/gravitino/api/expressions/sorts/null_ordering.py new file mode 100644 index 00000000000..a65a6efc97b --- /dev/null +++ b/clients/client-python/gravitino/api/expressions/sorts/null_ordering.py @@ -0,0 +1,35 @@ +# 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 enum import Enum + + +class NullOrdering(Enum): + """A null order used in sorting expressions.""" + + NULLS_FIRST: str = "nulls_first" + """Nulls appear before non-nulls. For ascending order, this means nulls appear at the beginning.""" + + NULLS_LAST: str = "nulls_last" + """Nulls appear after non-nulls. For ascending order, this means nulls appear at the end.""" + + def __str__(self) -> str: + if self == NullOrdering.NULLS_FIRST: + return "nulls_first" + if self == NullOrdering.NULLS_LAST: + return "nulls_last" + + raise ValueError(f"Unexpected null order: {self}") diff --git a/clients/client-python/gravitino/api/expressions/sorts/sort_direction.py b/clients/client-python/gravitino/api/expressions/sorts/sort_direction.py new file mode 100644 index 00000000000..23694b019cb --- /dev/null +++ b/clients/client-python/gravitino/api/expressions/sorts/sort_direction.py @@ -0,0 +1,73 @@ +# 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 enum import Enum +from gravitino.api.expressions.sorts.null_ordering import NullOrdering + + +class SortDirection(Enum): + """A sort direction used in sorting expressions. + Each direction has a default null ordering that is implied if no null ordering is specified explicitly. + """ + + ASCENDING = ("asc", NullOrdering.NULLS_FIRST) + """Ascending sort direction. Nulls appear first. For ascending order, this means nulls appear at the beginning.""" + + DESCENDING = ("desc", NullOrdering.NULLS_LAST) + """Descending sort direction. Nulls appear last. For ascending order, this means nulls appear at the end.""" + + def __init__(self, direction: str, default_null_ordering: NullOrdering): + self._direction = direction + self._default_null_ordering = default_null_ordering + + def direction(self) -> str: + return self._direction + + def default_null_ordering(self) -> NullOrdering: + """ + Returns the default null ordering to use if no null ordering is specified explicitly. + + Returns: + NullOrdering: The default null ordering. + """ + return self._default_null_ordering + + def __str__(self) -> str: + if self == SortDirection.ASCENDING: + return SortDirection.ASCENDING.direction() + if self == SortDirection.DESCENDING: + return SortDirection.DESCENDING.direction() + + raise ValueError(f"Unexpected sort direction: {self}") + + @staticmethod + def from_string(direction: str): + """ + Returns the SortDirection from the string representation. + + Args: + direction: The string representation of the sort direction. + + Returns: + SortDirection: The corresponding SortDirection. + """ + direction = direction.lower() + if direction == SortDirection.ASCENDING.direction(): + return SortDirection.ASCENDING + if direction == SortDirection.DESCENDING.direction(): + return SortDirection.DESCENDING + + raise ValueError(f"Unexpected sort direction: {direction}") diff --git a/clients/client-python/gravitino/api/expressions/sorts/sort_order.py b/clients/client-python/gravitino/api/expressions/sorts/sort_order.py new file mode 100644 index 00000000000..ae7a1bb27b2 --- /dev/null +++ b/clients/client-python/gravitino/api/expressions/sorts/sort_order.py @@ -0,0 +1,45 @@ +# 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 abc import ABC, abstractmethod +from typing import List + +from gravitino.api.expressions.expression import Expression +from gravitino.api.expressions.sorts.null_ordering import NullOrdering +from gravitino.api.expressions.sorts.sort_direction import SortDirection + + +class SortOrder(Expression, ABC): + """Represents a sort order in the public expression API.""" + + @abstractmethod + def expression(self) -> Expression: + """Returns the sort expression.""" + pass + + @abstractmethod + def direction(self) -> SortDirection: + """Returns the sort direction.""" + pass + + @abstractmethod + def null_ordering(self) -> NullOrdering: + """Returns the null ordering.""" + pass + + def children(self) -> List[Expression]: + """Returns the children expressions of this sort order.""" + return [self.expression()] diff --git a/clients/client-python/gravitino/api/expressions/sorts/sort_orders.py b/clients/client-python/gravitino/api/expressions/sorts/sort_orders.py new file mode 100644 index 00000000000..9deaa4bacd9 --- /dev/null +++ b/clients/client-python/gravitino/api/expressions/sorts/sort_orders.py @@ -0,0 +1,94 @@ +# 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 typing import List + +from gravitino.api.expressions.expression import Expression +from gravitino.api.expressions.sorts.null_ordering import NullOrdering +from gravitino.api.expressions.sorts.sort_direction import SortDirection +from gravitino.api.expressions.sorts.sort_order import SortOrder + + +class SortImpl(SortOrder): + + def __init__( + self, + expression: Expression, + direction: SortDirection, + null_ordering: NullOrdering, + ): + """Initialize the SortImpl object.""" + self._expression = expression + self._direction = direction + self._null_ordering = null_ordering + + def expression(self) -> Expression: + return self._expression + + def direction(self) -> SortDirection: + return self._direction + + def null_ordering(self) -> NullOrdering: + return self._null_ordering + + def __eq__(self, other: object) -> bool: + """Check if two SortImpl instances are equal.""" + if not isinstance(other, SortImpl): + return False + return ( + self.expression() == other.expression() + and self.direction() == other.direction() + and self.null_ordering() == other.null_ordering() + ) + + def __hash__(self) -> int: + """Generate a hash for a SortImpl instance.""" + return hash((self.expression(), self.direction(), self.null_ordering())) + + def __str__(self) -> str: + """Provide a string representation of the SortImpl object.""" + return ( + f"SortImpl(expression={self._expression}, " + f"direction={self._direction}, null_ordering={self._null_ordering})" + ) + + +class SortOrders: + """Helper methods to create SortOrders to pass into Apache Gravitino.""" + + # NONE is used to indicate that there is no sort order + NONE: List[SortOrder] = [] + + @staticmethod + def ascending(expression: Expression) -> SortImpl: + """Creates a sort order with ascending direction and nulls first.""" + return SortOrders.of(expression, SortDirection.ASCENDING) + + @staticmethod + def descending(expression: Expression) -> SortImpl: + """Creates a sort order with descending direction and nulls last.""" + return SortOrders.of(expression, SortDirection.DESCENDING) + + @staticmethod + def of( + expression: Expression, + direction: SortDirection, + null_ordering: NullOrdering = None, + ) -> SortImpl: + """Creates a sort order with the given direction and optionally specified null ordering.""" + if null_ordering is None: + null_ordering = direction.default_null_ordering() + return SortImpl(expression, direction, null_ordering) diff --git a/clients/client-python/tests/unittests/rel/test_sorts.py b/clients/client-python/tests/unittests/rel/test_sorts.py new file mode 100644 index 00000000000..7116add4de5 --- /dev/null +++ b/clients/client-python/tests/unittests/rel/test_sorts.py @@ -0,0 +1,118 @@ +# 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 unittest +from unittest.mock import MagicMock + +from gravitino.api.expressions.function_expression import FunctionExpression +from gravitino.api.expressions.named_reference import NamedReference +from gravitino.api.expressions.sorts.sort_direction import SortDirection +from gravitino.api.expressions.sorts.null_ordering import NullOrdering +from gravitino.api.expressions.sorts.sort_orders import SortImpl, SortOrders +from gravitino.api.expressions.expression import Expression + + +class TestSortOrder(unittest.TestCase): + def test_sort_direction_from_string(self): + self.assertEqual(SortDirection.from_string("asc"), SortDirection.ASCENDING) + self.assertEqual(SortDirection.from_string("desc"), SortDirection.DESCENDING) + with self.assertRaises(ValueError): + SortDirection.from_string("invalid") + + def test_null_ordering(self): + self.assertEqual(str(NullOrdering.NULLS_FIRST), "nulls_first") + self.assertEqual(str(NullOrdering.NULLS_LAST), "nulls_last") + + def test_sort_impl_initialization(self): + mock_expression = MagicMock(spec=Expression) + sort_impl = SortImpl( + expression=mock_expression, + direction=SortDirection.ASCENDING, + null_ordering=NullOrdering.NULLS_FIRST, + ) + self.assertEqual(sort_impl.expression(), mock_expression) + self.assertEqual(sort_impl.direction(), SortDirection.ASCENDING) + self.assertEqual(sort_impl.null_ordering(), NullOrdering.NULLS_FIRST) + + def test_sort_impl_equality(self): + mock_expression1 = MagicMock(spec=Expression) + mock_expression2 = MagicMock(spec=Expression) + + sort_impl1 = SortImpl( + expression=mock_expression1, + direction=SortDirection.ASCENDING, + null_ordering=NullOrdering.NULLS_FIRST, + ) + sort_impl2 = SortImpl( + expression=mock_expression1, + direction=SortDirection.ASCENDING, + null_ordering=NullOrdering.NULLS_FIRST, + ) + sort_impl3 = SortImpl( + expression=mock_expression2, + direction=SortDirection.ASCENDING, + null_ordering=NullOrdering.NULLS_FIRST, + ) + + self.assertEqual(sort_impl1, sort_impl2) + self.assertNotEqual(sort_impl1, sort_impl3) + + def test_sort_orders(self): + mock_expression = MagicMock(spec=Expression) + ascending_order = SortOrders.ascending(mock_expression) + self.assertEqual(ascending_order.direction(), SortDirection.ASCENDING) + self.assertEqual(ascending_order.null_ordering(), NullOrdering.NULLS_FIRST) + + descending_order = SortOrders.descending(mock_expression) + self.assertEqual(descending_order.direction(), SortDirection.DESCENDING) + self.assertEqual(descending_order.null_ordering(), NullOrdering.NULLS_LAST) + + def test_sort_impl_string_representation(self): + mock_expression = MagicMock(spec=Expression) + sort_impl = SortImpl( + expression=mock_expression, + direction=SortDirection.ASCENDING, + null_ordering=NullOrdering.NULLS_FIRST, + ) + expected_str = ( + f"SortImpl(expression={mock_expression}, " + f"direction=asc, null_ordering=nulls_first)" + ) + self.assertEqual(str(sort_impl), expected_str) + + def test_sort_order(self): + field_reference = NamedReference.field(["field1"]) + sort_order = SortOrders.of( + field_reference, SortDirection.ASCENDING, NullOrdering.NULLS_FIRST + ) + + self.assertEqual(NullOrdering.NULLS_FIRST, sort_order.null_ordering()) + self.assertEqual(SortDirection.ASCENDING, sort_order.direction()) + self.assertIsInstance(sort_order.expression(), NamedReference) + self.assertEqual(["field1"], sort_order.expression().field_name()) + + date = FunctionExpression.of("date", NamedReference.field(["b"])) + sort_order = SortOrders.of( + date, SortDirection.DESCENDING, NullOrdering.NULLS_LAST + ) + self.assertEqual(NullOrdering.NULLS_LAST, sort_order.null_ordering()) + self.assertEqual(SortDirection.DESCENDING, sort_order.direction()) + + self.assertIsInstance(sort_order.expression(), FunctionExpression) + self.assertEqual("date", sort_order.expression().function_name()) + self.assertEqual( + ["b"], sort_order.expression().arguments()[0].references()[0].field_name() + )