commands

=== Data Loader 1 ===


import io
import pandas as pd
import requests
if 'data_loader' not in globals():
    from mage_ai.data_preparation.decorators import data_loader
if 'test' not in globals():
    from mage_ai.data_preparation.decorators import test


@data_loader
def load_data_from_api(*args, **kwargs):
    """
    Template for loading data from API
    """
    url = 'https://storage.googleapis.com/new-york-cab-data/yellow_tripdata_2022_sample.csv'
    response = requests.get(url)

    return pd.read_csv(io.StringIO(response.text), sep=',')


@test
def test_output(output, *args) -> None:
    """
    Template code for testing the output of the block.
    """
    assert output is not None, 'The output is undefined'


# === Data Loader 2 ===


import io
import pandas as pd
import requests
if 'data_loader' not in globals():
    from mage_ai.data_preparation.decorators import data_loader
if 'test' not in globals():
    from mage_ai.data_preparation.decorators import test


@data_loader
def load_data_from_api(*args, **kwargs):
    """
    Template for loading data from API
    """
    url = 'https://storage.googleapis.com/new-york-cab-data/taxi_zones.csv'
    response = requests.get(url)

    return pd.read_csv(io.StringIO(response.text), sep=',')


@test
def test_output(output, *args) -> None:
    """
    Template code for testing the output of the block.
    """
    assert output is not None, 'The output is undefined'


=== Transformer 1 ===


import numpy as np
if 'transformer' not in globals():
    from mage_ai.data_preparation.decorators import transformer
if 'test' not in globals():
    from mage_ai.data_preparation.decorators import test


@transformer
def transform(df_taxi_zone, df, *args, **kwargs):
    """
    Template code for a transformer block.

    Add more parameters to this function if this block has multiple parent blocks.
    There should be one parameter for each output variable from each parent block.

    Args:
        data: The output from the upstream parent block
        args: The output from any additional upstream blocks (if applicable)

    Returns:
        Anything (e.g. data frame, dictionary, array, int, str, etc.)
    """
    # Specify your transformation logic here

    def get_coordinates(id, column='x'):
        try:
            return df_taxi_zone[df_taxi_zone['location_id'] == id][[f'coordinate_{column}']].values[0][0]
        except IndexError:
            return np.NaN

    df.dropna(inplace=True)
    print('drop data')
    df["PULocation_x"]= df.apply(lambda row: get_coordinates(row["PULocationID"]), axis=1)
    print('PULocation_x done')
    df["PULocation_y"]= df.apply(lambda row: get_coordinates(row["PULocationID"], 'y'), axis=1)
    print('PULocation_y done')
    df["DOLocation_x"]= df.apply(lambda row: get_coordinates(row["DOLocationID"]), axis=1)
    print('DOLocation_x done')
    df["DOLocation_y"]= df.apply(lambda row: get_coordinates(row["PULocationID"],'y'), axis=1)
    print('DOLocation_y done')

    df.drop(['PULocationID', 'DOLocationID'], axis=1, inplace=True)
    print('drop cols')
    df.dropna(inplace=True)


    return df


@test
def test_output(output, *args) -> None:
    """
    Template code for testing the output of the block.
    """
    assert output is not None, 'The output is undefined'


=== Transformer 2 ===


import pandas as pd

if 'transformer' not in globals():
    from mage_ai.data_preparation.decorators import transformer
if 'test' not in globals():
    from mage_ai.data_preparation.decorators import test


@transformer
def transform(df, *args, **kwargs):
    """
    Template code for a transformer block.

    Add more parameters to this function if this block has multiple parent blocks.
    There should be one parameter for each output variable from each parent block.

    Args:
        data: The output from the upstream parent block
        args: The output from any additional upstream blocks (if applicable)

    Returns:
        Anything (e.g. data frame, dictionary, array, int, str, etc.)
    """
    # Specify your transformation logic here
    df['tpep_pickup_datetime'] = pd.to_datetime(df['tpep_pickup_datetime'])
    df['tpep_dropoff_datetime'] = pd.to_datetime(df['tpep_dropoff_datetime'])


    df = df.drop_duplicates().reset_index(drop=True)
    df['trip_id'] = df.index


    datetime_dim = df[['tpep_pickup_datetime', 'tpep_dropoff_datetime']].reset_index(drop=True)
    datetime_dim['pick_hour'] = datetime_dim['tpep_pickup_datetime'].dt.hour
    datetime_dim['pick_day'] = datetime_dim['tpep_pickup_datetime'].dt.day
    datetime_dim['pick_month'] = datetime_dim['tpep_pickup_datetime'].dt.month
    datetime_dim['pick_year'] = datetime_dim['tpep_pickup_datetime'].dt.year
    datetime_dim['pick_weekday'] = datetime_dim['tpep_pickup_datetime'].dt.weekday

    datetime_dim['drop_hour'] = datetime_dim['tpep_pickup_datetime'].dt.hour
    datetime_dim['drop_day'] = datetime_dim['tpep_pickup_datetime'].dt.day
    datetime_dim['drop_month'] = datetime_dim['tpep_pickup_datetime'].dt.month
    datetime_dim['drop_year'] = datetime_dim['tpep_pickup_datetime'].dt.year
    datetime_dim['drop_weekday'] = datetime_dim['tpep_pickup_datetime'].dt.weekday
    datetime_dim['datetime_id'] = datetime_dim.index
    datetime_dim = datetime_dim[['datetime_id', 'tpep_pickup_datetime', 'pick_hour', 'pick_day', 'pick_month', 'pick_year', 'pick_weekday', 'tpep_dropoff_datetime', 'drop_hour', 'drop_day', 'drop_month', 'drop_year', 'drop_weekday']]


    passenger_count_dim = df[['passenger_count']].reset_index(drop=True)
    passenger_count_dim['passenger_count_id'] = passenger_count_dim.index
    passenger_count_dim['passenger_count'] = passenger_count_dim['passenger_count'].astype('int')
    passenger_count_dim[['passenger_count_id', 'passenger_count']]


    trip_distance_dim = df[['trip_distance']].reset_index(drop=True)
    trip_distance_dim['trip_distance_id'] = trip_distance_dim.index
    trip_distance_dim[['trip_distance_id', 'trip_distance']]


    PU_location_dim = df[['PULocation_x', 'PULocation_y']].reset_index(drop=True)
    PU_location_dim['PU_location_id'] = PU_location_dim.index
    PU_location_dim[['PU_location_id', 'PULocation_x', 'PULocation_y']]


    DO_location_dim = df[['DOLocation_x', 'DOLocation_y']].reset_index(drop=True)
    DO_location_dim['DO_location_id'] = DO_location_dim.index
    DO_location_dim[['DO_location_id', 'DOLocation_x', 'DOLocation_y']]


    payment_type_dim = df[['payment_type']].reset_index(drop=True)
    payment_type_dim['payment_type_id'] = payment_type_dim.index
    payment_name_dict = {1: 'Credit card', 2: 'Cash', 3: 'No charge', 4: 'Dispute', 5: 'Unknown', 6: 'Voided trip'}
    payment_type_dim['payment_type_name'] = payment_type_dim['payment_type'].map(payment_name_dict)
    payment_type_dim[['payment_type_id', 'payment_type', 'payment_type_name']]


    rate_code_dim = df[['RatecodeID']].reset_index(drop=True)
    rate_code_dim['rate_code_id'] = rate_code_dim.index
    rate_code_name_dict = {1: 'Standard rate', 2: 'JFK', 3: 'Newark', 4: 'Nassau or Westchester', 5: 'Negotiated fare', 6: 'Group ride'}
    rate_code_dim['rate_code_name'] = rate_code_dim['RatecodeID'].map(rate_code_name_dict)
    rate_code_dim[['rate_code_id', 'RatecodeID', 'rate_code_name']]


    fact_table = df.merge(passenger_count_dim, left_on='trip_id', right_on='passenger_count_id').merge(trip_distance_dim, left_on='trip_id', right_on='trip_distance_id').merge(rate_code_dim, left_on='trip_id', right_on='rate_code_id').merge(PU_location_dim, left_on='trip_id', right_on='PU_location_id').merge(DO_location_dim, left_on='trip_id', right_on='DO_location_id').merge(datetime_dim, left_on='trip_id', right_on='datetime_id').merge(payment_type_dim, left_on='trip_id', right_on='payment_type_id')[['trip_id','VendorID', 'datetime_id', 'passenger_count_id', 'trip_distance_id', 'rate_code_id', 'store_and_fwd_flag', 'PU_location_id', 'DO_location_id', 'payment_type_id', 'fare_amount', 'extra', 'mta_tax', 'tip_amount', 'tolls_amount',  'improvement_surcharge', 'total_amount']]


    return_dict = {'datetime_dim': datetime_dim.to_dict(orient='dict'),
    'passenger_count_dim': passenger_count_dim.to_dict(orient='dict'),
    'trip_distance_dim': trip_distance_dim.to_dict(orient='dict'),
    'PU_location_dim': PU_location_dim.to_dict(orient='dict'),
    'DO_location_dim': DO_location_dim.to_dict(orient='dict'),
    'payment_type_dim': payment_type_dim.to_dict(orient='dict'),
    'rate_code_dim': rate_code_dim.to_dict(orient='dict'),
    'fact_table': fact_table.to_dict(orient='dict'),}

    return return_dict


@test
def test_output(output, *args) -> None:
    """
    Template code for testing the output of the block.
    """
    assert output is not None, 'The output is undefined'


=== Data Explorer ===


from mage_ai.data_preparation.repo_manager import get_repo_path
from mage_ai.io.bigquery import BigQuery
from mage_ai.io.config import ConfigFileLoader
from pandas import DataFrame
from os import path

if 'data_exporter' not in globals():
    from mage_ai.data_preparation.decorators import data_exporter


@data_exporter
def export_data_to_big_query(data, **kwargs) -> None:
    """
    Template for exporting data to a BigQuery warehouse.
    Specify your configuration settings in 'io_config.yaml'.

    Docs: https://docs.mage.ai/design/data-loading#bigquery
    """
    config_path = path.join(get_repo_path(), 'io_config.yaml')
    config_profile = 'default'

    for key, value in data.items():
        table_id = f'new-york-cab-data-analysis.nyc_taxi_data.{key}'

        BigQuery.with_config(ConfigFileLoader(config_path, config_profile)).export(
            DataFrame(value),
            table_id,
            if_exists='replace',  # Specify resolution policy if table name already exists
        )