abacusai 0.63.0__tar.gz → 0.63.1__tar.gz

{abacusai-0.63.0/abacusai.egg-info → abacusai-0.63.1}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 1.1
 Name: abacusai
-Version: 0.63.0
+Version: 0.63.1
 Summary: Abacus.AI Python Client Library
 Home-page: https://github.com/abacusai/api-python
 Author: Abacus.AI

{abacusai-0.63.0 → abacusai-0.63.1}/abacusai/api_client_utils.py

@@ -21,3 +21,27 @@ def get_clean_function_source_code(func: Callable):
     # If function source code has some initial indentation, remove it (Ex - can happen if the functor was defined inside a function)
     source_code = dedent(source_code)
     return source_code
+
+
+def avro_to_pandas_dtype(avro_type):
+    avro_pandas_dtypes = {
+        'long': 'Int64',
+        'int': 'Int32',
+        'float': 'float32',
+        'double': 'float64',
+        'string': 'object',
+        'boolean': 'bool',
+        'bytes': 'object',
+        'null': 'object',
+    }
+
+    if isinstance(avro_type, dict):
+        avro_type = avro_type['type']
+
+    return avro_pandas_dtypes.get(avro_type, 'object')
+
+
+def get_non_nullable_type(types):
+    non_nullable_types = [
+        avro_type for avro_type in types if avro_type != 'null']
+    return non_nullable_types[0] if non_nullable_types else None

{abacusai-0.63.0 → abacusai-0.63.1}/abacusai/client.py

@@ -637,7 +637,7 @@ class ReadOnlyClient(BaseApiClient):
             FeatureGroup: All the feature groups associated with the specified Python function ID."""
         return self._call_api('listPythonFunctionFeatureGroups', 'GET', query_params={'name': name, 'limit': limit}, parse_type=FeatureGroup)
 
-    def execute_async_fg_operation(self, query: str = None) -> ExecuteFgOperation:
+    def execute_async_feature_group_operation(self, query: str = None) -> ExecuteFgOperation:
         """Starts the execution of fg operation
 
         Args:
@@ -645,9 +645,9 @@ class ReadOnlyClient(BaseApiClient):
 
         Returns:
             ExecuteFgOperation: A dict that contains the execution status"""
-        return self._call_api('executeAsyncFGOperation', 'GET', query_params={'query': query}, parse_type=ExecuteFgOperation)
+        return self._call_api('executeAsyncFeatureGroupOperation', 'GET', query_params={'query': query}, parse_type=ExecuteFgOperation)
 
-    def download_execute_fg_operation_result_part_chunk(self, execute_fg_operation_run_id: str, part: int, offset: int = 0, chunk_size: int = 10485760) -> Dict:
+    def download_execute_feature_group_operation_result_part_chunk(self, execute_fg_operation_run_id: str, part: int, offset: int = 0, chunk_size: int = 10485760) -> Dict:
         """Downloads a chunk of the result of the execution of fg operation
 
         Args:
@@ -655,7 +655,7 @@ class ReadOnlyClient(BaseApiClient):
             part (int): The part number of the result
             offset (int): The offset in the part
             chunk_size (int): The size of the chunk"""
-        return self._call_api('downloadExecuteFGOperationResultPartChunk', 'GET', query_params={'executeFgOperationRunId': execute_fg_operation_run_id, 'part': part, 'offset': offset, 'chunkSize': chunk_size})
+        return self._call_api('downloadExecuteFeatureGroupOperationResultPartChunk', 'GET', query_params={'executeFgOperationRunId': execute_fg_operation_run_id, 'part': part, 'offset': offset, 'chunkSize': chunk_size})
 
     def get_feature_group_version_export_download_url(self, feature_group_export_id: str) -> FeatureGroupExportDownloadUrl:
         """Get a link to download the feature group version.
@@ -2452,7 +2452,7 @@ class ApiClient(ReadOnlyClient):
         agent_function_name = agent_function.__name__
         return self.update_agent(model_id=model_id, function_source_code=function_source_code, agent_function_name=agent_function_name, memory=memory, package_requirements=package_requirements)
 
-    def execute_fg_sql(self, sql):
+    def execute_feature_group_sql(self, sql):
         """
         Execute a SQL query on the feature groups
 
@@ -2462,7 +2462,7 @@ class ApiClient(ReadOnlyClient):
         Returns:
             pandas.DataFrame: The result of the query.
         """
-        execute_query = self.execute_async_fg_operation(sql)
+        execute_query = self.execute_async_feature_group_operation(sql)
         execute_query.wait_for_execution()
         return execute_query.load_as_pandas()
 
@@ -3236,7 +3236,7 @@ Creates a new feature group defined as the union of other feature group versions
             lookup_keys (list): List of feature names which can be used in the lookup API to restrict the computation to a set of dataset rows. These feature names have to correspond to underlying dataset columns."""
         return self._call_api('setFeatureGroupIndexingConfig', 'POST', query_params={}, body={'featureGroupId': feature_group_id, 'primaryKey': primary_key, 'updateTimestampKey': update_timestamp_key, 'lookupKeys': lookup_keys})
 
-    def describe_async_fg_operation(self, execute_fg_operation_run_id: str) -> ExecuteFgOperation:
+    def describe_async_feature_group_operation(self, execute_fg_operation_run_id: str) -> ExecuteFgOperation:
         """Gets the status of the execution of fg operation
 
         Args:
@@ -3244,14 +3244,14 @@ Creates a new feature group defined as the union of other feature group versions
 
         Returns:
             ExecuteFgOperation: A dict that contains the execution status"""
-        return self._call_api('describeAsyncFGOperation', 'POST', query_params={}, body={'executeFgOperationRunId': execute_fg_operation_run_id}, parse_type=ExecuteFgOperation)
+        return self._call_api('describeAsyncFeatureGroupOperation', 'POST', query_params={}, body={'executeFgOperationRunId': execute_fg_operation_run_id}, parse_type=ExecuteFgOperation)
 
-    def get_execute_fg_operation_result_part_count(self, execute_fg_operation_run_id: str) -> None:
+    def get_execute_feature_group_operation_result_part_count(self, execute_fg_operation_run_id: str) -> None:
         """Gets the number of parts in the result of the execution of fg operation
 
         Args:
             execute_fg_operation_run_id (str): The unique ID associated with the execution."""
-        return self._call_api('getExecuteFGOperationResultPartCount', 'POST', query_params={}, body={'executeFgOperationRunId': execute_fg_operation_run_id})
+        return self._call_api('getExecuteFeatureGroupOperationResultPartCount', 'POST', query_params={}, body={'executeFgOperationRunId': execute_fg_operation_run_id})
 
     def update_feature_group(self, feature_group_id: str, description: str = None) -> FeatureGroup:
         """Modify an existing Feature Group.

{abacusai-0.63.0 → abacusai-0.63.1}/abacusai/execute_fg_operation.py

@@ -72,7 +72,7 @@ class ExecuteFgOperation(AbstractApiClass):
         return self
 
     def describe(self):
-        return self.client.describe_async_fg_operation(self.execute_fg_operation_run_id)
+        return self.client.describe_async_feature_group_operation(self.execute_fg_operation_run_id)
 
     # internal call
     def _download_avro_file(self, file_part, tmp_dir):
@@ -80,7 +80,7 @@ class ExecuteFgOperation(AbstractApiClass):
         part_path = os.path.join(tmp_dir, f'{file_part}.avro')
         with open(part_path, 'wb') as file:
             while True:
-                with self.client._call_api('downloadExecuteFGOperationResultPartChunk', 'GET', query_params={'executeFgOperationRunId': self.execute_fg_operation_run_id, 'part': file_part, 'offset': offset}, streamable_response=True) as chunk:
+                with self.client._call_api('downloadExecuteFeatureGroupOperationResultPartChunk', 'GET', query_params={'executeFgOperationRunId': self.execute_fg_operation_run_id, 'part': file_part, 'offset': offset}, streamable_response=True) as chunk:
                     bytes_written = file.write(chunk.read())
                 if not bytes_written:
                     break
@@ -100,8 +100,8 @@ class ExecuteFgOperation(AbstractApiClass):
         import fastavro
         import pandas as pd
 
-        file_parts = range(self.client._call_api('getExecuteFGOperationResultPartCount', 'POST', query_params={
-                           'executeFgOperationRunId': self.execute_fg_operation_run_id}))
+        file_parts = range(self.client._call_api('getExecuteFeatureGroupOperationResultPartCount',
+                           'POST', query_params={'executeFgOperationRunId': self.execute_fg_operation_run_id}))
         data_df = pd.DataFrame()
         with ThreadPoolExecutor(max_workers=max_workers) as executor:
             with tempfile.TemporaryDirectory() as tmp_dir:

{abacusai-0.63.0 → abacusai-0.63.1}/abacusai/feature_group_version.py

@@ -222,6 +222,8 @@ class FeatureGroupVersion(AbstractApiClass):
         import fastavro
         import pandas as pd
 
+        from .api_client_utils import avro_to_pandas_dtype, get_non_nullable_type
+
         file_parts = range(self.client._call_api(
             '_getFeatureGroupVersionPartCount', 'GET', query_params={'featureGroupVersion': self.id}))
         data_df = pd.DataFrame()
@@ -234,9 +236,26 @@ class FeatureGroupVersion(AbstractApiClass):
                     part_path = future.result()
                     with open(part_path, 'rb') as part_data:
                         reader = fastavro.reader(part_data)
-                        df_parts.append(pd.DataFrame([r for r in reader]))
-                data_df = pd.concat(df_parts, ignore_index=True)
-            for col in data_df.columns:
-                if pd.core.dtypes.common.is_datetime64_ns_dtype(data_df[col]):
-                    data_df[col] = data_df[col].dt.tz_localize(None)
+                        schema = reader.writer_schema
+                        col_dtypes = {}
+                        for field in schema['fields']:
+                            field_name = field['name']
+                            field_type = field['type']
+                            if isinstance(field_type, list):
+                                field_type = get_non_nullable_type(field_type)
+                            pandas_dtype = avro_to_pandas_dtype(field_type)
+                            col_dtypes[field_name] = pandas_dtype
+                        df_part = pd.DataFrame.from_records(
+                            [r for r in reader], columns=col_dtypes.keys())
+
+                        for col in df_part.columns:
+                            if pd.core.dtypes.common.is_datetime64_ns_dtype(df_part[col]):
+                                df_part[col] = df_part[col].dt.tz_localize(
+                                    None)
+                            else:
+                                df_part[col] = df_part[col].astype(
+                                    col_dtypes[col])
+                        df_parts.append(df_part)
+            data_df = pd.concat(df_parts, ignore_index=True)
+
         return data_df

{abacusai-0.63.0 → abacusai-0.63.1/abacusai.egg-info}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 1.1
 Name: abacusai
-Version: 0.63.0
+Version: 0.63.1
 Summary: Abacus.AI Python Client Library
 Home-page: https://github.com/abacusai/api-python
 Author: Abacus.AI

{abacusai-0.63.0 → abacusai-0.63.1}/setup.py

@@ -1,7 +1,7 @@
 from setuptools import find_packages, setup
 
 
-__version__ = '0.63.0'
+__version__ = '0.63.1'
 
 
 setup(name='abacusai',