Store Module Enhancements: Update CSV, Pandas, PyArrow, and native Python store implementations with tests

Author: neomatamune

cosmotech/coal/store/__init__.py

@@ -0,0 +1,42 @@
+# Copyright (C) - 2023 - 2025 - Cosmo Tech
+# This document and all information contained herein is the exclusive property -
+# including all intellectual property rights pertaining thereto - of Cosmo Tech.
+# Any use, reproduction, translation, broadcasting, transmission, distribution,
+# etc., to any person is prohibited unless it has been previously and
+# specifically authorized by written means by Cosmo Tech.
+
+"""
+Store module.
+
+This module provides functions for working with the Store,
+including loading and converting data.
+"""
+
+# Re-export the Store class
+from cosmotech.coal.store.store import Store
+
+# Re-export functions from the csv module
+from cosmotech.coal.store.csv import (
+    store_csv_file,
+    convert_store_table_to_csv,
+)
+
+# Re-export functions from the native_python module
+from cosmotech.coal.store.native_python import (
+    store_pylist,
+    convert_table_as_pylist,
+)
+
+# Re-export functions from the pandas module (if available)
+
+from cosmotech.coal.store.pandas import (
+    store_dataframe,
+    convert_store_table_to_dataframe as convert_store_table_to_pandas_dataframe,
+)
+
+# Re-export functions from the pyarrow module (if available)
+
+from cosmotech.coal.store.pyarrow import (
+    store_table,
+    convert_store_table_to_dataframe as convert_store_table_to_pyarrow_table,
+)

cosmotech/coal/store/csv.py

@@ -1,3 +1,10 @@
+# Copyright (C) - 2023 - 2025 - Cosmo Tech
+# This document and all information contained herein is the exclusive property -
+# including all intellectual property rights pertaining thereto - of Cosmo Tech.
+# Any use, reproduction, translation, broadcasting, transmission, distribution,
+# etc., to any person is prohibited unless it has been previously and
+# specifically authorized by written means by Cosmo Tech.
+
 import pathlib
 
 import pyarrow.csv as pc
@@ -9,7 +16,7 @@ def store_csv_file(
     table_name: str,
     csv_path: pathlib.Path,
     replace_existsing_file: bool = False,
-    store=Store()
+    store=Store(),
 ):
     if not csv_path.exists():
         raise FileNotFoundError(f"File {csv_path} does not exists")
@@ -18,16 +25,14 @@ def store_csv_file(
     _c = data.column_names
     data = data.rename_columns([Store.sanitize_column(_column) for _column in _c])
 
-    store.add_table(table_name=table_name,
-                    data=data,
-                    replace=replace_existsing_file)
+    store.add_table(table_name=table_name, data=data, replace=replace_existsing_file)
 
 
 def convert_store_table_to_csv(
     table_name: str,
     csv_path: pathlib.Path,
     replace_existsing_file: bool = False,
-    store=Store()
+    store=Store(),
 ):
     if csv_path.name.endswith(".csv") and csv_path.exists() and not replace_existsing_file:
         raise FileExistsError(f"File {csv_path} already exists")

cosmotech/coal/store/native_python.py

@@ -1,3 +1,10 @@
+# Copyright (C) - 2023 - 2025 - Cosmo Tech
+# This document and all information contained herein is the exclusive property -
+# including all intellectual property rights pertaining thereto - of Cosmo Tech.
+# Any use, reproduction, translation, broadcasting, transmission, distribution,
+# etc., to any person is prohibited unless it has been previously and
+# specifically authorized by written means by Cosmo Tech.
+
 import pyarrow as pa
 
 from cosmotech.coal.store.store import Store
@@ -6,19 +13,13 @@
 def store_pylist(
     table_name: str,
     data: list[dict],
-    replace_existsing_file:
-    bool = False,
-    store=Store()
+    replace_existsing_file: bool = False,
+    store=Store(),
 ):
     data = pa.Table.from_pylist(data)
 
-    store.add_table(table_name=table_name,
-                    data=data,
-                    replace=replace_existsing_file)
+    store.add_table(table_name=table_name, data=data, replace=replace_existsing_file)
 
 
-def convert_table_as_pylist(
-    table_name: str,
-    store=Store()
-):
+def convert_table_as_pylist(table_name: str, store=Store()):
     return store.get_table(table_name).to_pylist()

cosmotech/coal/store/pandas.py

@@ -1,31 +1,26 @@
+# Copyright (C) - 2023 - 2025 - Cosmo Tech
+# This document and all information contained herein is the exclusive property -
+# including all intellectual property rights pertaining thereto - of Cosmo Tech.
+# Any use, reproduction, translation, broadcasting, transmission, distribution,
+# etc., to any person is prohibited unless it has been previously and
+# specifically authorized by written means by Cosmo Tech.
+
 import pyarrow
 
 from cosmotech.coal.store.store import Store
+import pandas as pd
 
-try:
-    import pandas as pd
-
-
-    def store_dataframe(
-        table_name: str,
-        dataframe: pd.DataFrame,
-        replace_existsing_file: bool = False,
-        store=Store()
-    ):
-
-        data = pyarrow.Table.from_pandas(dataframe)
-
-        store.add_table(table_name=table_name,
-                        data=data,
-                        replace=replace_existsing_file)
 
+def store_dataframe(
+    table_name: str,
+    dataframe: pd.DataFrame,
+    replace_existsing_file: bool = False,
+    store=Store(),
+):
+    data = pyarrow.Table.from_pandas(dataframe)
 
-    def convert_store_table_to_dataframe(
-        table_name: str,
-        store=Store()
-    ) -> pd.DataFrame:
+    store.add_table(table_name=table_name, data=data, replace=replace_existsing_file)
 
-        return store.get_table(table_name).to_pandas()
 
-except ModuleNotFoundError:
-    pass
+def convert_store_table_to_dataframe(table_name: str, store=Store()) -> pd.DataFrame:
+    return store.get_table(table_name).to_pandas()

cosmotech/coal/store/pyarrow.py

@@ -1,23 +1,23 @@
-from cosmotech.coal.store.store import Store
-
-try:
-    import pyarrow as pa
+# Copyright (C) - 2023 - 2025 - Cosmo Tech
+# This document and all information contained herein is the exclusive property -
+# including all intellectual property rights pertaining thereto - of Cosmo Tech.
+# Any use, reproduction, translation, broadcasting, transmission, distribution,
+# etc., to any person is prohibited unless it has been previously and
+# specifically authorized by written means by Cosmo Tech.
 
+from cosmotech.coal.store.store import Store
 
-    def store_table(
-        table_name: str,
-        data: pa.Table,
-        replace_existsing_file: bool = False,
-        store=Store()
-    ):
+import pyarrow as pa
 
-        store.add_table(table_name=table_name,
-                        data=data,
-                        replace=replace_existsing_file)
 
+def store_table(
+    table_name: str,
+    data: pa.Table,
+    replace_existsing_file: bool = False,
+    store=Store(),
+):
+    store.add_table(table_name=table_name, data=data, replace=replace_existsing_file)
 
-    def convert_store_table_to_dataframe(table_name: str, store=Store()) -> pa.Table:
-        return store.get_table(table_name)
 
-except ModuleNotFoundError:
-    pass
+def convert_store_table_to_dataframe(table_name: str, store=Store()) -> pa.Table:
+    return store.get_table(table_name)

cosmotech/coal/store/store.py

@@ -1,23 +1,29 @@
+# Copyright (C) - 2023 - 2025 - Cosmo Tech
+# This document and all information contained herein is the exclusive property -
+# including all intellectual property rights pertaining thereto - of Cosmo Tech.
+# Any use, reproduction, translation, broadcasting, transmission, distribution,
+# etc., to any person is prohibited unless it has been previously and
+# specifically authorized by written means by Cosmo Tech.
+
 import os
 import pathlib
 
 import pyarrow
 from adbc_driver_sqlite import dbapi
 
 from cosmotech.coal.utils.logger import LOGGER
+from cosmotech.orchestrator.utils.translate import T
 
 
 class Store:
-
     @staticmethod
     def sanitize_column(column_name: str) -> str:
         return column_name.replace(" ", "_")
 
     def __init__(
         self,
         reset=False,
-        store_location: pathlib.Path = pathlib.Path(os.environ.get("CSM_PARAMETERS_ABSOLUTE_PATH",
-                                                                   "."))
+        store_location: pathlib.Path = pathlib.Path(os.environ.get("CSM_PARAMETERS_ABSOLUTE_PATH", ".")),
     ):
         self.store_location = pathlib.Path(store_location) / ".coal/store"
         self.store_location.mkdir(parents=True, exist_ok=True)
@@ -33,23 +39,23 @@ def reset(self):
 
     def get_table(self, table_name: str) -> pyarrow.Table:
         if not self.table_exists(table_name):
-            raise ValueError(f"No table with name {table_name} exists")
+            raise ValueError(T("coal.errors.data.no_table").format(table_name=table_name))
         return self.execute_query(f"select * from {table_name}")
 
     def table_exists(self, table_name) -> bool:
         return table_name in self.list_tables()
 
     def get_table_schema(self, table_name: str) -> pyarrow.Schema:
         if not self.table_exists(table_name):
-            raise ValueError(f"No table with name {table_name} exists")
+            raise ValueError(T("coal.errors.data.no_table").format(table_name=table_name))
         with dbapi.connect(self._database) as conn:
             return conn.adbc_get_table_schema(table_name)
 
     def add_table(self, table_name: str, data=pyarrow.Table, replace: bool = False):
         with dbapi.connect(self._database, autocommit=True) as conn:
             with conn.cursor() as curs:
                 rows = curs.adbc_ingest(table_name, data, "replace" if replace else "create_append")
-                LOGGER.debug(f"Inserted {rows} rows in table {table_name}")
+                LOGGER.debug(T("coal.logs.data_transfer.rows_inserted").format(rows=rows, table_name=table_name))
 
     def execute_query(self, sql_query: str) -> pyarrow.Table:
         batch_size = 1024
@@ -58,7 +64,7 @@ def execute_query(self, sql_query: str) -> pyarrow.Table:
             try:
                 with dbapi.connect(self._database, autocommit=True) as conn:
                     with conn.cursor() as curs:
-                        curs.adbc_statement.set_options(**{"adbc.sqlite.query.batch_rows":str(batch_size)})
+                        curs.adbc_statement.set_options(**{"adbc.sqlite.query.batch_rows": str(batch_size)})
                         curs.execute(sql_query)
                         return curs.fetch_arrow_table()
             except OSError: