Flowfile 0.5.1__py3-none-any.whl → 0.5.4__py3-none-any.whl

flowfile_core/flowfile/flow_data_engine/utils.py

@@ -1,13 +1,15 @@
+import os
+from base64 import encodebytes
+from collections.abc import Callable, Iterable
+from typing import Any
+
 import polars as pl
-from flowfile_core.configs.settings import AVAILABLE_RAM, WORKER_URL
+import requests
+
 from flowfile_core.configs import logger
-from flowfile_core.flowfile.flow_data_engine.subprocess_operations import ExternalDfFetcher
-from flowfile_core.flowfile.flow_data_engine.subprocess_operations import Status
+from flowfile_core.configs.settings import AVAILABLE_RAM, WORKER_URL
+from flowfile_core.flowfile.flow_data_engine.subprocess_operations import ExternalDfFetcher, Status
 from flowfile_core.utils.utils import standardize_col_dtype
-import os
-from typing import List, Dict, Iterable, Callable, Any
-import requests
-from base64 import encodebytes
 
 
 def get_data_type(vals: Iterable[Any]):
@@ -15,25 +17,26 @@ def get_data_type(vals: Iterable[Any]):
     if len(types) == 1:
         return types.pop().__name__
     elif types == {float, int}:
-        return 'float'
+        return "float"
     else:
-        return 'str'
+        return "str"
 
 
-def calculate_schema(lf: pl.LazyFrame) -> List[Dict]:
-    r = ExternalDfFetcher(lf=lf, operation_type='calculate_schema', wait_on_completion=False, flow_id=-1, node_id=-1)
-    schema_stats: List[Dict] = r.get_result()
+def calculate_schema(lf: pl.LazyFrame) -> list[dict]:
+    r = ExternalDfFetcher(lf=lf, operation_type="calculate_schema", wait_on_completion=False, flow_id=-1, node_id=-1)
+    schema_stats: list[dict] = r.get_result()
 
     for schema_stat in schema_stats:
-        schema_stat['pl_datatype'] = getattr(pl.datatypes, schema_stat['pl_datatype'])
+        schema_stat["pl_datatype"] = getattr(pl.datatypes, schema_stat["pl_datatype"])
 
     return schema_stats
 
 
-def write_polars_frame(_df: pl.LazyFrame | pl.DataFrame, path: str, data_type: str = 'parquet',
-                       estimated_size: int = 0):
+def write_polars_frame(
+    _df: pl.LazyFrame | pl.DataFrame, path: str, data_type: str = "parquet", estimated_size: int = 0
+):
     is_lazy = isinstance(_df, pl.LazyFrame)
-    logger.info('Caching data frame')
+    logger.info("Caching data frame")
     if is_lazy:
         if estimated_size > 0:
             fit_memory = estimated_size / 1024 / 1000 / 1000 < AVAILABLE_RAM
@@ -43,16 +46,16 @@ def write_polars_frame(_df: pl.LazyFrame | pl.DataFrame, path: str, data_type: s
 
         if is_lazy:
             logger.info("Writing in memory efficient mode")
-            write_method = getattr(_df, 'sink_' + data_type)
+            write_method = getattr(_df, "sink_" + data_type)
             try:
                 write_method(path)
                 return True
-            except Exception as e:
+            except Exception:
                 pass
         if is_lazy:
             _df = _df.collect()
     try:
-        write_method = getattr(_df, 'write_' + data_type)
+        write_method = getattr(_df, "write_" + data_type)
         write_method(path)
         return True
     except:
@@ -67,13 +70,13 @@ def collect(df: pl.LazyFrame, streamable: bool = True):
 
 
 def cache_polars_frame_to_temp(_df: pl.LazyFrame | pl.DataFrame, tempdir: str = None) -> pl.LazyFrame:
-    path = f'{tempdir}\\fl_file_{id(_df)}'
+    path = f"{tempdir}\\fl_file_{id(_df)}"
     result = write_polars_frame(_df, path)
     if result:
         df = pl.read_parquet(path)
         return df.lazy()
     else:
-        raise Exception('Could not cache the data')
+        raise Exception("Could not cache the data")
 
 
 def define_pl_col_transformation(col_name: str, col_type: pl.DataType) -> pl.Expr:
@@ -85,64 +88,94 @@ def define_pl_col_transformation(col_name: str, col_type: pl.DataType) -> pl.Exp
         return pl.col(col_name).cast(col_type, strict=False)
 
 
-def execute_write_method(write_method: Callable, path: str, data_type: str = None, sheet_name: str = None,
-                         delimiter: str = None,
-                         write_mode: str = 'create'):
-    if data_type == 'excel':
-        logger.info('Writing as excel file')
+def execute_write_method(
+    write_method: Callable,
+    path: str,
+    data_type: str = None,
+    sheet_name: str = None,
+    delimiter: str = None,
+    write_mode: str = "create",
+):
+    if data_type == "excel":
+        logger.info("Writing as excel file")
         write_method(path, worksheet=sheet_name)
-    elif data_type == 'csv':
-        logger.info('Writing as csv file')
-        if write_mode == 'append':
-            with open(path, 'ab') as f:
-                write_method(path=f, separator=delimiter, quote_style='always')
+    elif data_type == "csv":
+        logger.info("Writing as csv file")
+        if write_mode == "append":
+            with open(path, "ab") as f:
+                write_method(path=f, separator=delimiter, quote_style="always")
         else:
-            write_method(path=path, separator=delimiter, quote_style='always')
-    elif data_type == 'parquet':
-        logger.info('Writing as parquet file')
+            write_method(path=path, separator=delimiter, quote_style="always")
+    elif data_type == "parquet":
+        logger.info("Writing as parquet file")
         write_method(path)
 
 
-def write_output(_df: pl.LazyFrame,
-                 data_type: str, path: str, write_mode: str, sheet_name: str = None,
-                 delimiter: str = None, flow_id: int = -1, node_id: int | str = -1) -> Status:
+def write_output(
+    _df: pl.LazyFrame,
+    data_type: str,
+    path: str,
+    write_mode: str,
+    sheet_name: str = None,
+    delimiter: str = None,
+    flow_id: int = -1,
+    node_id: int | str = -1,
+) -> Status:
     serializable_df = _df.serialize()
-    r = requests.post(f'{WORKER_URL}/write_results/',
-                      json={'operation': encodebytes(serializable_df).decode(),
-                            'data_type': data_type,
-                            'path': path,
-                            'write_mode': write_mode,
-                            'sheet_name': sheet_name,
-                            'delimiter': delimiter,
-                            'flowfile_node_id': node_id,
-                            'flowfile_flow_id': flow_id})
+    r = requests.post(
+        f"{WORKER_URL}/write_results/",
+        json={
+            "operation": encodebytes(serializable_df).decode(),
+            "data_type": data_type,
+            "path": path,
+            "write_mode": write_mode,
+            "sheet_name": sheet_name,
+            "delimiter": delimiter,
+            "flowfile_node_id": node_id,
+            "flowfile_flow_id": flow_id,
+        },
+    )
     if r.ok:
         return Status(**r.json())
     else:
-        raise Exception(f'Could not cache the data, {r.text}')
-
-
-def local_write_output(_df: pl.LazyFrame | pl.DataFrame, data_type: str, path: str, write_mode: str,
-                       sheet_name: str = None, delimiter: str = None, flow_id: int = -1, node_id: int | str = -1):
+        raise Exception(f"Could not cache the data, {r.text}")
+
+
+def local_write_output(
+    _df: pl.LazyFrame | pl.DataFrame,
+    data_type: str,
+    path: str,
+    write_mode: str,
+    sheet_name: str = None,
+    delimiter: str = None,
+    flow_id: int = -1,
+    node_id: int | str = -1,
+):
     is_lazy = isinstance(_df, pl.LazyFrame)
-    sink_method_str = 'sink_' + data_type
-    write_method_str = 'write_' + data_type
+    sink_method_str = "sink_" + data_type
+    write_method_str = "write_" + data_type
     has_sink_method = hasattr(_df, sink_method_str)
     write_method = None
-    if os.path.exists(path) and write_mode == 'create':
+    if os.path.exists(path) and write_mode == "create":
         return None
     if has_sink_method and is_lazy:
-        write_method = getattr(_df, 'sink_' + data_type)
+        write_method = getattr(_df, "sink_" + data_type)
     elif not is_lazy or not has_sink_method:
         if is_lazy:
             _df = _df.collect()
         write_method = getattr(_df, write_method_str)
     if write_method is not None:
-        execute_write_method(write_method, path=path, data_type=data_type, sheet_name=sheet_name,
-                             delimiter=delimiter, write_mode=write_mode)
+        execute_write_method(
+            write_method,
+            path=path,
+            data_type=data_type,
+            sheet_name=sheet_name,
+            delimiter=delimiter,
+            write_mode=write_mode,
+        )
 
 
-def create_pl_df_type_save(raw_data: Iterable[Iterable], orient: str = 'row') -> pl.DataFrame:
+def create_pl_df_type_save(raw_data: Iterable[Iterable], orient: str = "row") -> pl.DataFrame:
     """
         orient : {'col', 'row'}, default None
         Whether to interpret two-dimensional data as columns or as rows. If None,
@@ -152,25 +185,24 @@ def create_pl_df_type_save(raw_data: Iterable[Iterable], orient: str = 'row') ->
     :param orient:
     :return: polars dataframe
     """
-    if orient == 'row':
-        raw_data = zip(*raw_data)
+    if orient == "row":
+        raw_data = zip(*raw_data, strict=False)
     raw_data = [standardize_col_dtype(values) for values in raw_data]
-    return pl.DataFrame(raw_data, orient='col')
+    return pl.DataFrame(raw_data, orient="col")
 
 
-def find_first_positions(lst: List[str]) -> Dict[str, int]:
-    first_positions: Dict[str, int] = {}
+def find_first_positions(lst: list[str]) -> dict[str, int]:
+    first_positions: dict[str, int] = {}
     for i, value in enumerate(lst):
         if value not in first_positions:
             first_positions[value] = i
     return first_positions
 
 
-def match_order(l: List[str], ref: List[str]) -> List[str]:
+def match_order(l: list[str], ref: list[str]) -> list[str]:
     ref_order = find_first_positions(ref)
     order = []
     for v in l:
-        org_order = ref_order.get(v, float('inf'))
+        org_order = ref_order.get(v, float("inf"))
         order.append(org_order)
-    return [v for _, v in sorted(zip(order, l))]
-
+    return [v for _, v in sorted(zip(order, l, strict=False))]