<jupyter_start><jupyter_text>Type FlexibilityIn the previous section, we saw how the `transform()` function was used to bring a Python function to Spark or Dask. In this section, we take a deeper look at the `transform()` function and what types of inputs and outputs it can handle. In the last section, the function we brought to Spark looked like this:```def predict(df: pd.DataFrame, model: LinearRegression) -> pd.DataFrame:    return df.assign(predicted=model.predict(df))```It had an input of type `pd.DataFrame`, and an output of type `pd.DataFrame`. Fugue applies these type annotations and converts the input to the type specified. The output annotation is also used to convert back to the engine DataFrame (Pandas, Spark or Dask). In this section, we will see some other examples of type annotations Fugue can handle.For those more familiar with distributed computing, the conversion is done for each partition. This will be explained more in later sections, but the important concept is that Fugue **does not** convert the whole Spark or Dask DataFrame to Pandas to perform the operation. The conversion overhead has been [benchmarked](../appendix/fugue_spark_benchmark.ipynb) and found to be negligible, especially for big data.  Sample Problem and `transform()`In this sample problem, we are interested in getting the first three digits of the `phone` column and populating a new column called `location` by using a dictionary that maps the values. We start by preparing the sample data and defining the mapping.<jupyter_code>import pandas as pd

_area_code_map = {"217": "Champaign, IL", "407": "Orlando, FL", "510": "Fremont, CA"}

data = pd.DataFrame({"phone": ["(217)-123-4567", "(217)-234-5678", "(407)-123-4567", 
                               "(407)-234-5678", "(510)-123-4567"]})<jupyter_output><empty_output><jupyter_text>First, we'll perform the operation in pandas. It's very simple because of the `.map()` method in pandas. We then test the function.<jupyter_code>def map_phone_to_location(df: pd.DataFrame) -> pd.DataFrame:
    df["location"] = df["phone"].str.slice(1,4).map(_area_code_map)
    return df

map_phone_to_location(data.copy())<jupyter_output><empty_output><jupyter_text>Similar to the function in the previous section, this function can be used in Spark on big data with Fugue's `transform()` function. For now, we will leave the engine blank to use the default pandas-based `NativeExecutionEngine`.<jupyter_code>from fugue import transform

transform(data.copy(),
          map_phone_to_location,
          schema="*, location:str").head(2)<jupyter_output><empty_output><jupyter_text>Different Input and Output AnnotationsThe `map_phone_to_location()` logic can actually be expressed in native Python. Below are three valid implementations of the same function. The `map_phone_to_location3()` below is less practical, but it's just to demo the varying types Fugue can take. `List[List]` will be useful in some use cases where you want to perform row-wise operations on DataFrames.<jupyter_code>from typing import List, Dict, Any, Iterable

def map_phone_to_location2(df: List[Dict[str,Any]]) -> Iterable[Dict[str,Any]]:
    for row in df:
        row["location"] = _area_code_map[row["phone"][1:4]]
        yield row

def map_phone_to_location3(df: List[List[Any]]) -> List[List[Any]]:
    for row in df:
        row.append(_area_code_map[row[0][1:4]])
    return df

def map_phone_to_location4(df: List[List[Any]]) -> pd.DataFrame:
    for row in df:
        row.append(_area_code_map[row[0][1:4]])
    df = pd.DataFrame.from_records(df, columns=["phone", "location"])
    return df<jupyter_output><empty_output><jupyter_text>Note that `map_phone_to_location4()` shows that the input and output types can differ. We can test these functions by passing some input.<jupyter_code>map_phone_to_location2([{"phone": "(217)-123-4567"}, {"phone": "(217)-234-5678"}])
map_phone_to_location3([["(217)-123-4567"], ["(217)-234-5678"]])
map_phone_to_location4([["(217)-123-4567"], ["(217)-234-5678"]])<jupyter_output><empty_output><jupyter_text>Even if these functions are not meant to operate on pandas DataFrames, Fugue `transform()` will handle the conversion for us so we don't need to make any changes to them to use them on Pandas, Spark or Dask.<jupyter_code>transform(data.copy(),
          map_phone_to_location2,
          schema="*, location:str").head(2)
transform(data.copy(),
          map_phone_to_location3,
          schema="*, location:str").head(2)
transform(data.copy(),
          map_phone_to_location4,
          schema="*, location:str").head(2)<jupyter_output><empty_output><jupyter_text>Bring the Functions to Distrbited ComputeBecause they work on the `NativeExecutionEngine`, we can also use the functions on the `SparkExecutionEngine`<jupyter_code>import fugue_spark

spark_df = transform(data.copy(),
                     map_phone_to_location2,    # Iterable[List] implementation
                     schema="*, location:str",
                     engine="spark")

spark_df.show(2)<jupyter_output>+--------------+-------------+
|         phone|     location|
+--------------+-------------+
|(217)-123-4567|Champaign, IL|
|(217)-234-5678|Champaign, IL|
+--------------+-------------+
only showing top 2 rows
{'fraction_non_alphanumeric': 0.1249050873196659, 'fraction_numerical': 0.03170083523158694, 'mean_word_length': 4.354674796747967, 'pattern_counts': {'":': 6, '<': 28, '<?xml version=': 0, '>': 33, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '7', 'dir': 'jupyter-structured-clean-dedup', 'id': '133587', 'license': 'Apache-2.0', 'n_tokens_mistral': 1677, 'n_tokens_neox': 1533, 'n_words': 531, 'path': 'tutorials/beginner/type_flexibility.ipynb', 'repo_name': 'fugue-project/tutorials', 'size': '5268'}