# ---
# jupyter:
#   jupytext:
#     text_representation:
#       extension: .py
#       format_name: light
#       format_version: '1.5'
#       jupytext_version: 1.14.4
#   kernelspec:
#     display_name: conda_python3
#     language: python
#     name: conda_python3
# ---

import pandas as pd
import random, string
from time import gmtime, strftime, sleep
import datetime
import boto3
from sagemaker.session import Session
from sagemaker.feature_store.feature_group import FeatureGroup
from sagemaker import get_execution_role
import time


# +
def str_time_prop(start, end, format):
    stime = time.mktime(time.strptime(start, format))
    etime = time.mktime(time.strptime(end, format))
    ptime = stime + random.random() * (etime - stime)
    return time.strftime(format, time.localtime(ptime))

def random_date(start, end):
    return str_time_prop(start, end, "%Y-%m-%dT%H:%M:%SZ")


# +
df = pd.DataFrame()
random.seed(42)

for i in range(3):
    df_tmp = pd.read_csv('s3://sagemaker-sample-files/datasets/tabular/fraud_detection/synthethic_fraud_detection_SA/sampled_transactions.csv')
    
    if i == 0:
        start = "2021-01-01T20:00:00Z"
        end = "2021-01-02T10:00:00Z"
        s = list(range(0, 2000))
    elif i == 1:
        start = "2021-01-04T20:00:00Z"
        end = "2021-01-05T10:00:00Z"
        s = list(range(2000, 4000))
    elif i == 2:
        start = "2021-01-07T20:00:00Z"
        end = "2021-01-08T10:00:00Z"
        s = list(range(4000, 6000))
    
    df_tmp['EventTime'] = df_tmp.apply(lambda x: random_date(start, end), axis=1)
    df_tmp = df_tmp.set_index([s])
    df = df.append(df_tmp)
# -

df[df['TransactionID'] == 3343087]

idx = df.groupby(['TransactionID'])['EventTime'].transform(max) == df['EventTime']
df_online = df[idx].copy()
df_offline = df[~idx].copy()


def cast_object_to_string(data_frame):
    for label in data_frame.columns:
        if data_frame.dtypes[label] == 'object':
            data_frame[label] = data_frame[label].astype("str").astype("string")


cast_object_to_string(df_offline)

cast_object_to_string(df_online)

# +
role = get_execution_role()
region = boto3.Session().region_name
boto_session = boto3.Session(region_name=region)
sagemaker_client = boto_session.client(service_name='sagemaker', region_name=region)
featurestore_runtime = boto_session.client(service_name='sagemaker-featurestore-runtime', region_name=region)
featuregroup_name = 'transactions-fg-manual-ingest-online-sync'
account_id = boto3.client('sts').get_caller_identity()["Account"]

feature_store_session = Session(
    boto_session=boto_session,
    sagemaker_client=sagemaker_client,
    sagemaker_featurestore_runtime_client=featurestore_runtime
)
# -

feature_group = FeatureGroup(name=featuregroup_name, sagemaker_session=feature_store_session)
feature_group.load_feature_definitions(data_frame=df_offline)

# +
record_identifier_feature_name = "TransactionID"
event_time_feature_name = "EventTime"

bucket = feature_store_session.default_bucket()
s3_folder = 'feature-store-manual-ingestion-online-sync'

def wait_for_feature_group_creation_complete(feature_group):
    status = feature_group.describe().get("FeatureGroupStatus")
    while status == "Creating":
        print("Waiting for Feature Group Creation")
        sleep(5)
        status = feature_group.describe().get("FeatureGroupStatus")
    if status != "Created":
        raise RuntimeError(f"Failed to create feature group {feature_group.name}")
    print(f"FeatureGroup {feature_group.name} successfully created.")

feature_group.create(
    s3_uri=f"s3://{bucket}/{s3_folder}",
    record_identifier_name=record_identifier_feature_name,
    event_time_feature_name=event_time_feature_name,
    role_arn=role,
    enable_online_store=True
)

wait_for_feature_group_creation_complete(feature_group=feature_group)
# -

# ## Backfilling historical records into S3

query = feature_group.athena_query()
fg_table = query.table_name

df_offline['key'] = df_offline['EventTime'].apply(lambda x: datetime.datetime.strptime(x, "%Y-%m-%dT%H:%M:%SZ").strftime("%Y-%m-%d-%H"))

fg_timestamp = strftime("%Y-%m-%dT%H:%M:%SZ", gmtime())
df_offline['write_time'] = df_offline['api_invocation_time'] = pd.to_datetime(fg_timestamp)
df_offline['is_deleted'] = False


def create_s3_paths(df):
    df_copy = df.copy()
    df_copy['EventTime'] = df_copy['EventTime'].astype(object)
    s3_paths = {}
    
    # loop over the unique keys
    for v in df['key'].unique():
        year, month, day, hour = v.split('-')
        # create path
        path = f"s3://{bucket}/{s3_folder}/{account_id}/sagemaker/{region}/offline-store/{fg_table}/data/year={year}/month={month}/day={day}/hour={hour}/"
        # identify the last entry for each group and retrieve minute and second
        _, minute, second = df_copy[df_copy['key'] == v].groupby('key')['EventTime'].max().iloc[0].split(':')
        # create filename
        filename = f"{year}{month}{day}T{hour}{minute}{second}_"
        filename += ''.join(random.choice(string.ascii_uppercase + string.ascii_lowercase + string.digits) for _ in range(16))
        filename += '.parquet'
        # append path + filename to dictionary
        s3_paths[v] = path + filename
        
    return s3_paths


# create a key for each row in the data that contains year, month, day, hour
df_offline['key'] = df_offline['EventTime'].apply(lambda x: datetime.datetime.strptime(x, "%Y-%m-%dT%H:%M:%SZ").strftime("%Y-%m-%d-%H"))
# create the correspong S3 paths and filenames
file_path_names = create_s3_paths(df_offline)

for key, data in df_offline.groupby('key'):
    data = data.drop('key', axis=1)
    data.to_parquet(file_path_names[key])

# +
query_string = f'SELECT * FROM "{fg_table}"'

query.run(query_string=query_string, output_location=f's3://{bucket}/{s3_folder}/query_results/')
query.wait()
dataset = query.as_dataframe()

dataset.head()
# -

# ## Writing most recent records into Online Store

feature_group.ingest(
    data_frame=df_online, max_workers=10, wait=True
)

record_identifier_value = str(2990130)
featurestore_runtime.get_record(FeatureGroupName=featuregroup_name, RecordIdentifierValueAsString=record_identifier_value)

# Running the query below immediately will return 2 records from the offline store, because it takes the Ingest API call ~5 minutes to update the Offline store. After ~5 minutes the query should return 3 queries.

# +
query_string = f'SELECT * FROM "{fg_table}" WHERE transactionid = 2990130'

query.run(query_string=query_string, output_location=f's3://{bucket}/{s3_folder}/query_results/')
query.wait()
dataset = query.as_dataframe()

dataset.head()
# -

feature_group.delete()

s3 = boto3.resource('s3')
s3bucket = s3.Bucket(bucket)
s3bucket.objects.filter(Prefix=s3_folder).delete()

StarcoderdataJupyterScriptsDedupFiltered