diff --git a/sdk/cosmos/azure-cosmos/CHANGELOG.md b/sdk/cosmos/azure-cosmos/CHANGELOG.md
index 260bf6ce75c8..5c5a6b109a8a 100644
--- a/sdk/cosmos/azure-cosmos/CHANGELOG.md
+++ b/sdk/cosmos/azure-cosmos/CHANGELOG.md
@@ -3,8 +3,11 @@
 ### 4.3.0b4 (Unreleased)
 
 #### Features Added
-- Added support for AAD authentication for the async client
-- Added support for AAD authentication for the sync client
+- Added support for AAD authentication for the async client.
+- Added support for AAD authentication for the sync client.
+
+#### Other Changes
+- Changed `_set_partition_key` return typehint in async client.
 
 ### 4.3.0b3 (2022-03-10)
 
diff --git a/sdk/cosmos/azure-cosmos/azure/cosmos/aio/container.py b/sdk/cosmos/azure-cosmos/azure/cosmos/aio/container.py
index f66fc278de60..50c47c4026fd 100644
--- a/sdk/cosmos/azure-cosmos/azure/cosmos/aio/container.py
+++ b/sdk/cosmos/azure-cosmos/azure/cosmos/aio/container.py
@@ -22,7 +22,7 @@
 """Create, read, update and delete items in the Azure Cosmos DB SQL API service.
 """
 
-from typing import Any, Dict, List, Optional, Union, cast
+from typing import Any, Dict, List, Optional, Union, cast, Awaitable
 from azure.core.async_paging import AsyncItemPaged
 
 from azure.core.tracing.decorator import distributed_trace  # pylint: disable=unused-import
@@ -106,11 +106,12 @@ def _get_conflict_link(self, conflict_or_link):
             return u"{}/conflicts/{}".format(self.container_link, conflict_or_link)
         return conflict_or_link["_self"]
 
-    async def _set_partition_key(self, partition_key):
+    def _set_partition_key(self, partition_key) -> Union[str, Awaitable]:
         if partition_key == NonePartitionKeyValue:
-            return CosmosClientConnection._return_undefined_or_empty_partition_key(await self.is_system_key)
+            return CosmosClientConnection._return_undefined_or_empty_partition_key(self.is_system_key)
         return partition_key
 
+
     @distributed_trace_async
     async def read(
             self,
@@ -231,8 +232,7 @@ async def read_item(
         doc_link = self._get_document_link(item)
         request_options = _build_options(kwargs)
         response_hook = kwargs.pop('response_hook', None)
-        if partition_key is not None:
-            request_options["partitionKey"] = await self._set_partition_key(partition_key)
+        request_options["partitionKey"] = self._set_partition_key(partition_key)
         max_integrated_cache_staleness_in_ms = kwargs.pop('max_integrated_cache_staleness_in_ms', None)
         if max_integrated_cache_staleness_in_ms is not None:
             validate_cache_staleness_value(max_integrated_cache_staleness_in_ms)
@@ -399,7 +399,7 @@ def query_items_change_feed(
             feed_options["partitionKeyRangeId"] = partition_key_range_id
         partition_key = kwargs.pop("partitionKey", None)
         if partition_key is not None:
-            feed_options["partitionKey"] = partition_key
+            feed_options["partitionKey"] = self._set_partition_key(partition_key)
         if is_start_from_beginning is not None:
             feed_options["isStartFromBeginning"] = is_start_from_beginning
         if max_item_count is not None:
@@ -534,8 +534,7 @@ async def delete_item(
         """
         request_options = _build_options(kwargs)
         response_hook = kwargs.pop('response_hook', None)
-        if partition_key is not None:
-            request_options["partitionKey"] = await self._set_partition_key(partition_key)
+        request_options["partitionKey"] = self._set_partition_key(partition_key)
         if pre_trigger_include is not None:
             request_options["preTriggerInclude"] = pre_trigger_include
         if post_trigger_include is not None:
@@ -691,9 +690,7 @@ async def read_conflict(
         """
         request_options = _build_options(kwargs)
         response_hook = kwargs.pop('response_hook', None)
-        if partition_key is not None:
-            request_options["partitionKey"] = await self._set_partition_key(partition_key)
-
+        request_options["partitionKey"] = self._set_partition_key(partition_key)
         result = await self.client_connection.ReadConflict(
             conflict_link=self._get_conflict_link(conflict), options=request_options, **kwargs
         )
@@ -721,9 +718,7 @@ async def delete_conflict(
         """
         request_options = _build_options(kwargs)
         response_hook = kwargs.pop('response_hook', None)
-        if partition_key is not None:
-            request_options["partitionKey"] = await self._set_partition_key(partition_key)
-
+        request_options["partitionKey"] = self._set_partition_key(partition_key)
         result = await self.client_connection.DeleteConflict(
             conflict_link=self._get_conflict_link(conflict), options=request_options, **kwargs
         )
diff --git a/sdk/cosmos/azure-cosmos/azure/cosmos/partition_key.py b/sdk/cosmos/azure-cosmos/azure/cosmos/partition_key.py
index ed2ab5f8f167..9c2be4cd19dc 100644
--- a/sdk/cosmos/azure-cosmos/azure/cosmos/partition_key.py
+++ b/sdk/cosmos/azure-cosmos/azure/cosmos/partition_key.py
@@ -22,6 +22,7 @@
 """Create partition keys in the Azure Cosmos DB SQL API service.
 """
 
+
 class NonePartitionKeyValue(object):
     """Represents None value for partitionKey when it's missing in a container.
     """