From 6a09ff49d7fd817eee6bacc01c772e9f326db9ce Mon Sep 17 00:00:00 2001
From: Thomas Petit-Jean <30775613+Thomzoy@users.noreply.github.com>
Date: Tue, 19 Mar 2024 17:10:36 +0100
Subject: [PATCH] fix: cast single date columns to datetime

---
 changelog.md          |  3 ++-
 eds_scikit/io/hive.py | 33 ++++++++++++++++++++++++++++++++-
 2 files changed, 34 insertions(+), 2 deletions(-)

diff --git a/changelog.md b/changelog.md
index 613ad5b3..03b591c1 100644
--- a/changelog.md
+++ b/changelog.md
@@ -3,7 +3,8 @@
 ## Unreleased
 
 ### Fixed
-- pyarrow fix did not work on spark executors
+- Pyarrow fix now work on spark executors.
+- Fix OMOP _date columns issue
 
 ## v0.1.7 (2024-04-12)
 ### Changed
diff --git a/eds_scikit/io/hive.py b/eds_scikit/io/hive.py
index f04dd378..54c42106 100644
--- a/eds_scikit/io/hive.py
+++ b/eds_scikit/io/hive.py
@@ -6,12 +6,16 @@
 
 import pandas as pd
 import pyarrow.parquet as pq
+import pyspark.sql.functions as F
+import pyspark.sql.types as T
 from databricks import koalas
 from loguru import logger
 from pyspark.sql import DataFrame as SparkDataFrame
 from pyspark.sql import SparkSession
 from pyspark.sql.types import LongType, StructField, StructType
 
+from eds_scikit.utils.framework import cache
+
 from . import settings
 from .base import BaseData
 from .data_quality import clean_dates
@@ -33,6 +37,8 @@ def __init__(
             Union[Dict[str, Optional[List[str]]], List[str]]
         ] = None,
         database_type: Optional[str] = "OMOP",
+        prune_omop_date_columns: bool = True,
+        cache: bool = True,
     ):
         """Spark interface for OMOP data stored in a Hive database.
 
@@ -54,6 +60,12 @@ def __init__(
             *deprecated*
         database_type: Optional[str] = 'OMOP'. Must be 'OMOP' or 'I2B2'
             Whether to use the native OMOP schema or to convert I2B2 inputs to OMOP.
+        prune_omop_date_columns: bool, default=True
+            In OMOP, most date values are stored both in a `<str>_date` and `<str>_datetime` column
+            Koalas has trouble handling the `date` time, so we only keep the `datetime` column
+        cache: bool, default=True
+            Whether to cache each table after preprocessing or not.
+            Will speed-up subsequent calculations, but can be long/infeasable for very large tables
 
         Attributes
         ----------
@@ -125,6 +137,8 @@ def __init__(
             for omop_table, i2b2_table in self.omop_to_i2b2.items():
                 self.i2b2_to_omop[i2b2_table].append(omop_table)
 
+        self.prune_omop_date_columns = prune_omop_date_columns
+        self.cache = cache
         self.user = os.environ["USER"]
         self.person_ids, self.person_ids_df = self._prepare_person_ids(person_ids)
         self.available_tables = self.list_available_tables()
@@ -224,10 +238,27 @@ def _read_table(
         if "person_id" in df.columns and person_ids is not None:
             df = df.join(person_ids, on="person_id", how="inner")
 
-        df = df.cache().to_koalas()
+        if self.prune_omop_date_columns:
+
+            # Keeping only _datetime column if corresponding _date exists
+            cols = [
+                c
+                for c in df.columns
+                if not ((c.endswith("_date") and (f"{c}time" in df.columns)))
+            ]
+            df = df.select(cols)
+
+            # Casting the single _date columns to timestamp:
+            for col in df.schema:
+                if col.dataType == T.DateType():
+                    df = df.withColumn(col.name, F.col(col.name).cast("timestamp"))
+        df = df.to_koalas()
 
         df = clean_dates(df)
 
+        if self.cache:
+            df = cache(df)
+
         return df
 
     def persist_tables_to_folder(