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
@@ -27,11 +34,7 @@ def convert_store_table_to_csv(
     replace_existsing_file: bool = False,
     store=Store(),
 ):
-    if (
-        csv_path.name.endswith(".csv")
-        and csv_path.exists()
-        and not replace_existsing_file
-    ):
+    if csv_path.name.endswith(".csv") and csv_path.exists() and not replace_existsing_file:
         raise FileExistsError(f"File {csv_path} already exists")
     if not csv_path.name.endswith(".csv"):
         csv_path = csv_path / f"{table_name}.csv"

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

cosmotech/coal/store/pandas.py

@@ -1,26 +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)
+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
-        )
+    store.add_table(table_name=table_name, data=data, replace=replace_existsing_file)
 
-    def convert_store_table_to_dataframe(
-        table_name: str, store=Store()
-    ) -> pd.DataFrame:
-        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,20 +1,23 @@
+# 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
 
-try:
-    import pyarrow as pa
+import pyarrow as pa
+
 
-    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 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,10 +1,18 @@
+# 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:
@@ -15,9 +23,7 @@ def sanitize_column(column_name: str) -> str:
     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,25 +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}")
+                rows = curs.adbc_ingest(table_name, data, "replace" if replace else "create_append")
+                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
@@ -60,9 +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: