fixed snowpark and pandas dataframe type error

This commit is contained in:
CIEF ACC1
2023-11-02 09:50:30 +00:00
parent ea8317d5a3
commit 30a08de7a0
+28 -27
View File
@@ -17,7 +17,7 @@ from datetime import datetime as dt, timedelta
# Function to check if a given date is a weekend [Saturday(5) or Sunday(6)]
def is_weekend(date):
return (date.weekday() == 5) | (date.weekday() == 6)
return date.weekday() in (5, 6)
# Function to compute the in advance working days
@@ -28,14 +28,14 @@ def compute_working_days(start_date, num_working_days, holiday_list, after_cut_o
counter = 0
# For dates on weekend or holiday or after cut off time, additional 1 working day to the loop
if (after_cut_off == 1) | is_weekend(start_date) | (start_date in holiday_list):
if (after_cut_off == 1) or is_weekend(start_date) or (start_date in holiday_list):
counter -= 1
# Loop to increase n number of working days, if weekend/holiday, skip counter
while counter < num_working_days:
start_date = start_date + timedelta(days=1)
if is_weekend(start_date) | (start_date in holiday_list):
if is_weekend(start_date) or (start_date in holiday_list):
continue
counter += 1
@@ -45,36 +45,37 @@ def compute_working_days(start_date, num_working_days, holiday_list, after_cut_o
# Main function
def model(dbt, session):
# Setting configuration
dbt.config(materialized="table",
packages = ["pandas"])
# Import data from upstream dbt model
df_date = dbt.ref("int__dates")
sp_df_date = dbt.ref("int__dates")
# Extract holiday dates into a list
df_work = df_date.to_pandas()
df_new = df_work[df_work["IS_COMPANY_HOLIDAY"] == 1]
df_holiday_list = df_new["DATE_DAY"].tolist()
# Filter holiday dates into a dataframe
sp_df_filter = sp_df_date.filter(sp_df_date['IS_COMPANY_HOLIDAY'] == 1)
sp_df_holiday_date = sp_df_filter.select('DATE_DAY')
# Apply function to df
df_work["ADD_1_WORKING_DAY_INCLUDED_BASE_DATE"] = df_work.apply(lambda row: compute_working_days(row["DATE_DAY"], 1, df_holiday_list, 0), axis=1)
df_work["ADD_2_WORKING_DAY_INCLUDED_BASE_DATE"] = df_work.apply(lambda row: compute_working_days(row["DATE_DAY"], 2, df_holiday_list, 0), axis=1)
df_work["ADD_3_WORKING_DAY_INCLUDED_BASE_DATE"] = df_work.apply(lambda row: compute_working_days(row["DATE_DAY"], 3, df_holiday_list, 0), axis=1)
df_work["ADD_7_WORKING_DAY_INCLUDED_BASE_DATE"] = df_work.apply(lambda row: compute_working_days(row["DATE_DAY"], 7, df_holiday_list, 0), axis=1)
df_work["ADD_30_WORKING_DAY_INCLUDED_BASE_DATE"] = df_work.apply(lambda row: compute_working_days(row["DATE_DAY"], 30, df_holiday_list, 0), axis=1)
df_work["ADD_90_WORKING_DAY_INCLUDED_BASE_DATE"] = df_work.apply(lambda row: compute_working_days(row["DATE_DAY"], 90, df_holiday_list, 0), axis=1)
df_work["ADD_365_WORKING_DAY_INCLUDED_BASE_DATE"] = df_work.apply(lambda row: compute_working_days(row["DATE_DAY"], 365, df_holiday_list, 0), axis=1)
# Convert snowpark dataframe to pandas dataframe
pd_df_date = sp_df_date.to_pandas()
pd_df_holiday_date = sp_df_holiday_date.to_pandas()
# Store holiday dates in a list
holiday_list = pd_df_holiday_date['DATE_DAY'].tolist()
# For orders after cut off time 4pm
# Affects only the weekdays
df_work["ADD_1_WORKING_DAY_EXCLUDED_BASE_DATE"] = df_work.apply(lambda row: compute_working_days(row["DATE_DAY"], 1, df_holiday_list, 1), axis=1)
df_work["ADD_2_WORKING_DAY_EXCLUDED_BASE_DATE"] = df_work.apply(lambda row: compute_working_days(row["DATE_DAY"], 2, df_holiday_list, 1), axis=1)
df_work["ADD_3_WORKING_DAY_EXCLUDED_BASE_DATE"] = df_work.apply(lambda row: compute_working_days(row["DATE_DAY"], 3, df_holiday_list, 1), axis=1)
df_work["ADD_7_WORKING_DAY_EXCLUDED_BASE_DATE"] = df_work.apply(lambda row: compute_working_days(row["DATE_DAY"], 7, df_holiday_list, 1), axis=1)
df_work["ADD_30_WORKING_DAY_EXCLUDED_BASE_DATE"] = df_work.apply(lambda row: compute_working_days(row["DATE_DAY"], 30, df_holiday_list, 1), axis=1)
df_work["ADD_90_WORKING_DAY_EXCLUDED_BASE_DATE"] = df_work.apply(lambda row: compute_working_days(row["DATE_DAY"], 90, df_holiday_list, 1), axis=1)
df_work["ADD_365_WORKING_DAY_EXCLUDED_BASE_DATE"] = df_work.apply(lambda row: compute_working_days(row["DATE_DAY"], 365, df_holiday_list, 1), axis=1)
# Applying add working days function to df
add_working_days = [1, 2, 3, 7, 30, 90 ,365]
return df_work
for working_day in add_working_days:
for after_cut_off in (0, 1):
if after_cut_off == 0:
col_name = f'ADD_{working_day}_WORKING_DAY_INCLUDED_BASE_DATE'
else:
col_name = f'ADD_{working_day}_WORKING_DAY_EXCLUDED_BASE_DATE'
# Apply function and append new calculated columns in the dataframe
pd_df_date[col_name] = pd_df_date['DATE_DAY'].apply(lambda start_date: compute_working_days(start_date, working_day, holiday_list, after_cut_off))
return pd_df_date