order by changes

Author: yashlimbad

e6data_python_connector/e6data_grpc.py

@@ -14,6 +14,7 @@
 from decimal import Decimal
 from io import BytesIO
 from ssl import CERT_NONE, CERT_OPTIONAL, CERT_REQUIRED
+from typing import overload
 
 import grpc
 from grpc._channel import _InactiveRpcError
@@ -710,6 +711,12 @@ def __init__(self, connection: Connection, file_path):
         self._batch = None
         self._create_dataframe()
 
+    def __enter__(self):
+        pass
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        pass
+
     def _create_dataframe(self):
         client = self.connection.client
 
@@ -726,7 +733,7 @@ def _create_dataframe(self):
         )
         self._query_id = create_dataframe_response.queryId
 
-    def select(self, *fields):
+    def select(self, *fields) -> "DataFrame":
         projection_fields = []
         for field in fields:
             projection_fields.append(field)
@@ -744,7 +751,7 @@ def select(self, *fields):
 
         return self
 
-    def where(self, where_clause : str):
+    def where(self, where_clause : str) -> "DataFrame":
         client = self.connection.client
         filter_on_dataframe_request = e6x_engine_pb2.FilterOnDataFrameRequest(
             queryId=self._query_id,
@@ -756,11 +763,14 @@ def where(self, where_clause : str):
             filter_on_dataframe_request
         )
 
-    def order_by(self, fields_list : list, sort_direction_list = None, null_direction_list = None):
+        return self
+
+    @overload
+    def order_by(self, field_list : list, sort_direction_list = None, null_direction_list = None) -> "DataFrame":
         orderby_fields = []
         sort_direction_request = []
         null_direction_request = []
-        for field in fields_list:
+        for field in field_list:
             orderby_fields.append(field)
 
         for direction in sort_direction_list:
@@ -796,7 +806,29 @@ def order_by(self, fields_list : list, sort_direction_list = None, null_directio
         )
         return self
 
-    def limit(self, fetch_limit : int):
+    def order_by(self, *field_list) -> "DataFrame":
+        orderby_fields = []
+        sort_direction_request = []
+        null_direction_request = []
+        for field in field_list:
+            orderby_fields.append(field)
+
+        client = self.connection.client
+
+        orderby_on_dataframe_request = e6x_engine_pb2.OrderByOnDataFrameRequest(
+            queryId=self._query_id,
+            sessionId=self._sessionId,
+            field=orderby_fields,
+            sortDirection=sort_direction_request,
+            nullsDirection=null_direction_request
+        )
+
+        orderby_on_dataframe_response = client.orderByOnDataFrame(
+            orderby_on_dataframe_request
+        )
+        return self
+
+    def limit(self, fetch_limit : int) -> "DataFrame":
         client = self.connection.client
         limit_on_dataframe_request = e6x_engine_pb2.LimitOnDataFrameRequest(
             queryId=self._query_id,
@@ -808,6 +840,8 @@ def limit(self, fetch_limit : int):
             limit_on_dataframe_request
         )
 
+        return self
+
     def show(self):
         self.execute()
         return self.fetchall()