diff --git a/sdk/python/feast/pyspark/__init__.py b/sdk/python/feast/pyspark/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/sdk/python/feast/pyspark/historical_feature_retrieval_job.py b/sdk/python/feast/pyspark/historical_feature_retrieval_job.py new file mode 100644 index 00000000000..d9afd43306d --- /dev/null +++ b/sdk/python/feast/pyspark/historical_feature_retrieval_job.py @@ -0,0 +1,421 @@ +import json +from datetime import timedelta +from typing import Any, Dict, List + +from pyspark import SparkFiles +from pyspark.sql import DataFrame, SparkSession, Window +from pyspark.sql.functions import col, expr, monotonically_increasing_id, row_number + + +def as_of_join( + entity: DataFrame, + entity_keys: List[str], + feature_table: DataFrame, + features: List[str], + feature_prefix: str = "", + max_age: int = None, +) -> DataFrame: + """Perform an as of join between entity and feature table, given a maximum age tolerance. + Join conditions: + 1. Entity primary key(s) value matches. + 2. Feature event timestamp is the closest match possible to the entity event timestamp, + but must not be more recent than the entity event timestamp, and the difference must + not be greater than max_age, unless max_age is not specified. + 3. If more than one feature table rows satisfy condition 1 and 2, feature row with the + most recent created timestamp will be chosen. + 4. If none of the above conditions are satisfied, the feature rows will have null values. + + Args: + entity (DataFrame): + Entity dataframe. Must contain the column event_timestamp. + entity_keys (List[str]): + Primary keys for the entity. + feature_table (DataFrame): + Feature table dataframe. Must contain the columns event_timestamp and created_timestamp. + features (List[str]): + The feature columns which should be present in the result dataframe. + feature_prefix (str): + Feature column prefix for the result dataframe. Useful for cases where the entity dataframe + contains one or more columns that share the same name as the features. + max_age (int): + Tolerance for the feature event timestamp recency, in seconds. + + Returns: + DataFrame: Join result. + + Example: + >>> entity.show() + +------+-------------------+ + |entity| event_timestamp| + +------+-------------------+ + | 1001|2020-09-02 00:00:00| + +------+-------------------+ + + >>> feature_table.show() + +------+-------+-------------------+-------------------+ + |entity|feature| event_timestamp| created_timestamp| + +------+-------+-------------------+-------------------+ + | 10| 200|2020-09-01 00:00:00|2020-09-02 00:00:00| + +------+-------+-------------------+-------------------+ + | 10| 400|2020-09-01 00:00:00|2020-09-01 00:00:00| + +------+-------+-------------------+-------------------+ + >>> df = as_of_join(entity, ["entity"], feature_table, ["feature"], feature_prefix = "prefix_") + >>> df.show() + +------+-------------------+--------------+ + |entity| event_timestamp|prefix_feature| + +------+-------------------+--------------+ + | 1001|2020-09-02 00:00:00| 200| + +------+-------------------+--------------+ + + >>> df = as_of_join(entity, ["entity"], feature_table, ["feature"], max_age = 12 * 60 * 60) + >>> df.show() + +------+-------------------+-------+ + |entity| event_timestamp|feature| + +------+-------------------+-------+ + | 1001|2020-09-02 00:00:00| null| + +------+-------------------+-------+ + + """ + entity_with_id = entity.withColumn("_row_nr", monotonically_increasing_id()) + + feature_event_timestamp = f"{feature_prefix}event_timestamp" + feature_created_timestamp = f"{feature_prefix}created_timestamp" + + projection = [ + col(col_name).alias(f"{feature_prefix}{col_name}") + for col_name in entity_keys + + features + + ["event_timestamp", "created_timestamp"] + ] + + selected_feature_table = feature_table.select(projection) + + join_cond = ( + entity_with_id.event_timestamp + >= selected_feature_table[feature_event_timestamp] + ) + if max_age: + join_cond = join_cond & ( + selected_feature_table[feature_event_timestamp] + >= entity_with_id.event_timestamp - expr(f"INTERVAL {max_age} seconds") + ) + + for key in entity_keys: + join_cond = join_cond & ( + entity_with_id[key] == selected_feature_table[f"{feature_prefix}{key}"] + ) + + conditional_join = entity_with_id.join( + selected_feature_table, join_cond, "leftOuter" + ) + for key in entity_keys: + conditional_join = conditional_join.drop( + selected_feature_table[f"{feature_prefix}{key}"] + ) + + window = Window.partitionBy("_row_nr", *entity_keys).orderBy( + col(feature_event_timestamp).desc(), col(feature_created_timestamp).desc() + ) + filter_most_recent_feature_timestamp = conditional_join.withColumn( + "_rank", row_number().over(window) + ).filter(col("_rank") == 1) + + return filter_most_recent_feature_timestamp.select( + entity.columns + [f"{feature_prefix}{feature}" for feature in features] + ) + + +class SchemaMismatchError(Exception): + pass + + +class MissingColumnError(Exception): + pass + + +class TimestampColumnError(Exception): + pass + + +def verify_schema( + df: DataFrame, expected_dtypes: Dict[str, str], is_feature_table: bool = False +): + """Verify if a dataframe has correct data types for as of joins. + Verification criteria: + 1. There is a column named `event_timestamp` of `Timestamp` type. + 2. For feature table, `created_timestamp` must be present as well. + 3. The dataframe should contains all the columns specified in `expected_dtypes`, + and has the same data type. + + Args: + df (DataFrame): + Input dataframe. + expected_dtypes (Dict[str, str]): + A map which defines the expected data types. The key is the column that + must be present in the dataframe, whereas the corresponding value is the + expected data types of that column. + is_feature_table (bool): + Whether the input dataframe represents a feature table. + + Raises: + MissingColumnError: If one or more required columns are missing. + TimestampColumnError: If the timestamp column has the wrong type. + SchemaMismatchError: If one or more column has the wrong data types as + compared to `expected_dtypes`""" + + actual_dtypes = dict(df.dtypes) + + if not set(expected_dtypes.keys()).issubset(set(actual_dtypes.keys())): + raise MissingColumnError("") + + for column, expected_type in expected_dtypes.items(): + if actual_dtypes[column] != expected_type: + raise SchemaMismatchError( + f"Schema mismatch. Expected schema: {expected_dtypes}, Actual schema: {actual_dtypes}", + ) + expected_timestamp_cols = ( + ["event_timestamp", "created_timestamp"] + if is_feature_table + else ["event_timestamp"] + ) + + for timestamp_col in expected_timestamp_cols: + if timestamp_col not in df.columns: + raise MissingColumnError( + f"{timestamp_col} not found. Input columns: {', '.join(df.columns)}" + ) + elif actual_dtypes[timestamp_col] != "timestamp": + raise TimestampColumnError( + f"{timestamp_col} is expected to be of timestamp type. Actual type: {actual_dtypes[timestamp_col]}" + ) + + +def join_entity_to_feature_tables( + query_conf: List[Dict[str, Any]], entity: DataFrame, tables: Dict[str, DataFrame] +) -> DataFrame: + """Perform as of join between entity and multiple feature table. Returns a DataFrame. + + Args: + query_conf (List[Dict[str, Any]]): + Query configuration. + entity (DataFrame): + Entity dataframe. Must contain the column event_timestamp. + tables (Dict[str, DataFrame]): + Map of feature table name to Spark DataFrame. + + Returns: + DataFrame: Join result. + + Example: + >>> entity.show() + +------+-------------------+ + |entity| event_timestamp| + +------+-------------------+ + | 1001|2020-09-02 00:00:00| + +------+-------------------+ + + >>> feature1.show() + +------+--------+-------------------+-------------------+ + |entity|feature1| event_timestamp| created_timestamp| + +------+--------+-------------------+-------------------+ + | 10| 200|2020-09-01 00:00:00|2020-09-01 00:00:00| + +------+--------+-------------------+------------------- + + >>> feature2.show() + +------+--------+-------------------+-------------------+ + |entity|feature2| event_timestamp| created_timestamp| + +------+--------+-------------------+-------------------+ + | 10| 400|2020-09-01 00:00:00|2020-09-01 00:00:00| + +------+--------+-------------------+------------------- + + + >>> tables = {"table1": feature1, "table2": feature2} + + >>> query_conf = [ + { + "table": "table1", + "features": ["feature1"], + "join": ["entity"], + }, + { + "table": "table2", + "features": ["feature2"], + "join": ["entity"], + }, + ] + + >>> joined_df = join_entity_to_feature_tables( + query_conf, + entity, + tables + ) + + >>> joined_df.show() + +------+-------------------+----------------+----------------+ + |entity| event_timestamp|table1__feature1|table2__feature2| + +------+-------------------+----------------+----------------+ + | 1001|2020-09-02 00:00:00| 200| 400| + +------+-------------------+----------------+----------------+ + """ + joined = entity + for query in query_conf: + joined = as_of_join( + joined, + query["join"], + tables[query["table"]], + query["features"], + feature_prefix=f"{query['table']}__", + max_age=query.get("max_age"), + ) + return joined + + +def retrieve_historical_features(spark: SparkSession, conf: Dict) -> DataFrame: + """Retrieve batch features based on given configuration. + + Args: + spark (SparkSession): + Spark session. + conf (Dict): + Configuration for the retrieval job, in json format. Sample configuration as follows: + + Returns: + DataFrame: Join result. + + Example: + sample_conf = { + "entity": { + "format": "csv", + "path": "file:///some_dir/customer_driver_pairs.csv", + "options": {"inferSchema": "true", "header": "true"}, + "col_mapping": { + "id": "driver_id" + }, + "dtypes": { + "driver_id": "integer" + } + }, + "tables": [ + { + "format": "parquet", + "path": "gs://some_bucket/bookings.parquet", + "name": "bookings", + "col_mapping": { + "id": "driver_id" + }, + "dtypes": { + "driver_id": "integer" + } + }, + { + "format": "avro", + "path": ""s3://some_bucket/transactions.parquet"", + "name": "transactions", + }, + ], + "queries": [ + { + "table": "transactions", + "features": ["daily_transactions"], + "join": ["customer_id"], + "max_age": 172800, + }, + { + "table": "bookings", + "features": ["completed_bookings"], + "join": ["driver_id"], + }, + ], + "output": + "format": "parquet" + "path": "gs://some_bucket/output.parquet" + } + + The values for the `format` and `path` should be recognizable by the Spark cluster where the job + is going to run on. For example, if you specify `bigquery` as input format, then you should ensure + that the Spark Big Query connector is installed on the cluster. Like wise, s3a connector is required + for Amazon S3 path. + + `options` is optional. If present, the options will be used when reading / writing the input / output. + + `max_age` is in seconds, and determines the lower bound of the timestamp of the retrieved feature. + If not specified, this would be unbounded. + + If necessary, `col_mapping` can be provided to map the columns of the dataframes before performing + the join operation. `col_mapping` is a dictionary where the key is the source column and the value + is the mapped column. + + `dtypes` is an optional parameter which helps the spark job to check whether the input dataframes + (after the col mapping) have the correct data types. The key is the column name, whereas the values + is the simple string format of the Spark data types, similar to the value returned by DataFrame.dtypes(). + Please refer to https://spark.apache.org/docs/latest/api/python/_modules/pyspark/sql/types.html for more + information. + + """ + + def map_column(df: DataFrame, col_mapping: Dict[str, str]): + projection = [ + col(col_name).alias(col_mapping.get(col_name, col_name)) + for col_name in df.columns + ] + return df.select(projection) + + entity = conf["entity"] + entity_df = ( + spark.read.format(entity["format"]) + .options(**entity.get("options", {})) + .load(entity["path"]) + ) + + entity_col_mapping = conf["entity"].get("col_mapping", {}) + mapped_entity_df = map_column(entity_df, entity_col_mapping) + verify_schema(mapped_entity_df, entity.get("dtypes", {})) + + tables = { + table_spec["name"]: map_column( + spark.read.format(table_spec["format"]) + .options(**table_spec.get("options", {})) + .load(table_spec["path"]), + table_spec.get("col_mapping", {}), + ) + for table_spec in conf["tables"] + } + + for table_spec in conf["tables"]: + verify_schema( + tables[table_spec["name"]], + table_spec.get("dtypes", {}), + is_feature_table=True, + ) + + max_timestamp = mapped_entity_df.agg({"event_timestamp": "max"}).collect()[0][0] + min_timestamp = mapped_entity_df.agg({"event_timestamp": "min"}).collect()[0][0] + + for query in conf["queries"]: + max_age = query.get("max_age") + if max_age: + tables[query["table"]] = tables[query["table"]].filter( + col("event_timestamp").between( + min_timestamp - timedelta(seconds=max_age), max_timestamp + ) + ) + + return join_entity_to_feature_tables(conf["queries"], mapped_entity_df, tables) + + +def start_job(spark: SparkSession, conf: Dict): + result = retrieve_historical_features(spark, conf) + output = conf["output"] + result.write.format(output["format"]).options(**output.get("options", {})).mode( + "overwrite" + ).save(output["path"]) + + +if __name__ == "__main__": + spark = SparkSession.builder.appName("Batch Retrieval").getOrCreate() + spark.sparkContext.addFile("config.json") + config_file_path = SparkFiles.get("config.json") + with open(config_file_path, "r") as config_file: + conf = json.load(config_file) + start_job(spark, conf) + spark.stop() diff --git a/sdk/python/requirements-ci.txt b/sdk/python/requirements-ci.txt index 20486cd1266..69c0be5a05e 100644 --- a/sdk/python/requirements-ci.txt +++ b/sdk/python/requirements-ci.txt @@ -9,6 +9,7 @@ pytest-lazy-fixture==0.6.3 pytest-mock pytest-timeout pytest-ordering==0.6.* +pyspark==3.* pandas~=1.0.0 mock==2.0.0 pandavro==1.5.* diff --git a/sdk/python/requirements-dev.txt b/sdk/python/requirements-dev.txt index ff2fbd812f0..8912392ce63 100644 --- a/sdk/python/requirements-dev.txt +++ b/sdk/python/requirements-dev.txt @@ -38,3 +38,5 @@ flake8 black==19.10b0 boto3 moto +pyspark==3.* +pyspark-stubs==3.* diff --git a/sdk/python/tests/data/bookings.csv b/sdk/python/tests/data/bookings.csv new file mode 100644 index 00000000000..b312cb896ef --- /dev/null +++ b/sdk/python/tests/data/bookings.csv @@ -0,0 +1,6 @@ +driver_id,event_timestamp,created_timestamp,completed_bookings +8001,2020-08-31T00:00:00.000,2020-08-31T00:00:00.000,200 +8001,2020-09-01T00:00:00.000,2020-09-01T00:00:00.000,300 +8002,2020-09-01T00:00:00.000,2020-09-01T00:00:00.000,600 +8002,2020-09-01T00:00:00.000,2020-09-02T00:00:00.000,500 +8003,2020-09-01T00:00:00.000,2020-09-02T00:00:00.000,700 diff --git a/sdk/python/tests/data/column_mapping_test_entity.csv b/sdk/python/tests/data/column_mapping_test_entity.csv new file mode 100644 index 00000000000..ec4b64c55c9 --- /dev/null +++ b/sdk/python/tests/data/column_mapping_test_entity.csv @@ -0,0 +1,6 @@ +id,event_timestamp +1001,2020-09-02T00:00:00.000 +1001,2020-09-03T00:00:00.000 +2001,2020-09-04T00:00:00.000 +2001,2020-09-04T00:00:00.000 +3001,2020-09-04T00:00:00.000 diff --git a/sdk/python/tests/data/column_mapping_test_feature.csv b/sdk/python/tests/data/column_mapping_test_feature.csv new file mode 100644 index 00000000000..ec8ee340efd --- /dev/null +++ b/sdk/python/tests/data/column_mapping_test_feature.csv @@ -0,0 +1,6 @@ +customer_id,total_bookings,datetime,created_datetime +1001,200,2020-09-02T00:00:00.000,2020-09-02T00:00:00.000 +1001,400,2020-09-04T00:00:00.000,2020-09-02T00:00:00.000 +2001,500,2020-09-03T00:00:00.000,2020-09-01T00:00:00.000 +2001,600,2020-09-03T00:00:00.000,2020-09-02T00:00:00.000 +3001,700,2020-09-03T00:00:00.000,2020-09-03T00:00:00.000 diff --git a/sdk/python/tests/data/customer_driver_pairs.csv b/sdk/python/tests/data/customer_driver_pairs.csv new file mode 100644 index 00000000000..32edf1db999 --- /dev/null +++ b/sdk/python/tests/data/customer_driver_pairs.csv @@ -0,0 +1,6 @@ +customer_id,driver_id,event_timestamp +1001,8001,2020-09-02T00:00:00.000 +1001,8002,2020-09-02T00:00:00.000 +1001,8002,2020-09-03T00:00:00.000 +2001,8002,2020-09-03T00:00:00.000 +2001,8002,2020-09-04T00:00:00.000 diff --git a/sdk/python/tests/data/transactions.csv b/sdk/python/tests/data/transactions.csv new file mode 100644 index 00000000000..977023b3a74 --- /dev/null +++ b/sdk/python/tests/data/transactions.csv @@ -0,0 +1,6 @@ +customer_id,event_timestamp,created_timestamp,daily_transactions +1001,2020-08-31T00:00:00.000,2020-09-01T00:00:00.000,50.0 +1001,2020-09-01T00:00:00.000,2020-09-01T00:00:00.000,100.0 +2001,2020-09-01T00:00:00.000,2020-08-31T00:00:00.000,80.0 +2001,2020-09-01T00:00:00.000,2020-09-01T00:00:00.000,200.0 +3001,2020-09-01T00:00:00.000,2020-09-01T00:00:00.000,300.0 diff --git a/sdk/python/tests/test_as_of_join.py b/sdk/python/tests/test_as_of_join.py new file mode 100644 index 00000000000..21cd829f29e --- /dev/null +++ b/sdk/python/tests/test_as_of_join.py @@ -0,0 +1,783 @@ +import os +import pathlib +import shutil +import tempfile +from datetime import datetime, timedelta +from os import path +from typing import Any, Dict, List + +import pytest +from pyspark.sql import DataFrame, SparkSession +from pyspark.sql.functions import lit +from pyspark.sql.types import ( + FloatType, + IntegerType, + StructField, + StructType, + TimestampType, +) + +from feast.pyspark.historical_feature_retrieval_job import ( + MissingColumnError, + SchemaMismatchError, + TimestampColumnError, + as_of_join, + join_entity_to_feature_tables, + retrieve_historical_features, + verify_schema, +) + + +@pytest.yield_fixture(scope="module") +def spark(pytestconfig): + spark_session = ( + SparkSession.builder.appName("Batch Retrieval Test") + .master("local") + .getOrCreate() + ) + yield spark_session + spark_session.stop() + + +@pytest.yield_fixture(scope="module") +def large_entity_csv_file(pytestconfig, spark): + start_datetime = datetime(year=2020, month=8, day=31) + nr_rows = 1000 + entity_data = [ + (1000 + i, start_datetime + timedelta(days=i)) for i in range(nr_rows) + ] + temp_dir = tempfile.mkdtemp() + file_path = os.path.join(temp_dir, "large_entity") + entity_schema = StructType( + [ + StructField("customer_id", IntegerType()), + StructField("event_timestamp", TimestampType()), + ] + ) + large_entity_df = spark.createDataFrame( + spark.sparkContext.parallelize(entity_data), entity_schema + ) + + large_entity_df.write.option("header", "true").csv(file_path) + yield file_path + shutil.rmtree(temp_dir) + + +@pytest.yield_fixture(scope="module") +def large_feature_csv_file(pytestconfig, spark): + start_datetime = datetime(year=2020, month=8, day=30) + nr_rows = 1000 + feature_data = [ + ( + 1000 + i, + start_datetime + timedelta(days=i), + start_datetime + timedelta(days=i + 1), + i * 10, + ) + for i in range(nr_rows) + ] + temp_dir = tempfile.mkdtemp() + file_path = os.path.join(temp_dir, "large_feature") + feature_schema = StructType( + [ + StructField("customer_id", IntegerType()), + StructField("event_timestamp", TimestampType()), + StructField("created_timestamp", TimestampType()), + StructField("total_bookings", IntegerType()), + ] + ) + large_feature_df = spark.createDataFrame( + spark.sparkContext.parallelize(feature_data), feature_schema + ) + + large_feature_df.write.option("header", "true").csv(file_path) + yield file_path + shutil.rmtree(temp_dir) + + +@pytest.fixture +def single_entity_schema(): + return StructType( + [ + StructField("customer_id", IntegerType()), + StructField("event_timestamp", TimestampType()), + ] + ) + + +@pytest.fixture +def composite_entity_schema(): + return StructType( + [ + StructField("customer_id", IntegerType()), + StructField("driver_id", IntegerType()), + StructField("event_timestamp", TimestampType()), + ] + ) + + +@pytest.fixture +def customer_feature_schema(): + return StructType( + [ + StructField("customer_id", IntegerType()), + StructField("event_timestamp", TimestampType()), + StructField("created_timestamp", TimestampType()), + StructField("daily_transactions", FloatType()), + ] + ) + + +@pytest.fixture +def driver_feature_schema(): + return StructType( + [ + StructField("driver_id", IntegerType()), + StructField("event_timestamp", TimestampType()), + StructField("created_timestamp", TimestampType()), + StructField("completed_bookings", IntegerType()), + ] + ) + + +@pytest.fixture +def rating_feature_schema(): + return StructType( + [ + StructField("customer_id", IntegerType()), + StructField("driver_id", IntegerType()), + StructField("event_timestamp", TimestampType()), + StructField("created_timestamp", TimestampType()), + StructField("customer_rating", FloatType()), + StructField("driver_rating", FloatType()), + ] + ) + + +def assert_dataframe_equal(left: DataFrame, right: DataFrame): + is_column_equal = set(left.columns) == set(right.columns) + + if not is_column_equal: + print(f"Column not equal. Left: {left.columns}, Right: {right.columns}") + assert is_column_equal + + is_content_equal = ( + left.exceptAll(right).count() == 0 and right.exceptAll(left).count() == 0 + ) + if not is_content_equal: + print("Rows are different.") + print("Left:") + left.show() + print("Right:") + right.show() + + assert is_content_equal + + +def test_join_without_max_age( + spark: SparkSession, + single_entity_schema: StructType, + customer_feature_schema: StructType, +): + entity_data = [ + (1001, datetime(year=2020, month=8, day=31)), + (1001, datetime(year=2020, month=9, day=1)), + (1001, datetime(year=2020, month=9, day=2)), + (1001, datetime(year=2020, month=9, day=3)), + (2001, datetime(year=2020, month=9, day=2)), + (3001, datetime(year=2020, month=9, day=1)), + ] + entity_df = spark.createDataFrame( + spark.sparkContext.parallelize(entity_data), single_entity_schema + ) + + feature_table_data = [ + ( + 1001, + datetime(year=2020, month=9, day=1), + datetime(year=2020, month=9, day=1), + 50.0, + ), + ( + 1001, + datetime(year=2020, month=9, day=1), + datetime(year=2020, month=9, day=2), + 100.0, + ), + ( + 2001, + datetime(year=2020, month=9, day=1), + datetime(year=2020, month=9, day=1), + 400.0, + ), + ( + 1001, + datetime(year=2020, month=9, day=2), + datetime(year=2020, month=9, day=1), + 200.0, + ), + ( + 1001, + datetime(year=2020, month=9, day=4), + datetime(year=2020, month=9, day=1), + 300.0, + ), + ] + feature_table_df = spark.createDataFrame( + spark.sparkContext.parallelize(feature_table_data), customer_feature_schema + ) + + joined_df = as_of_join( + entity_df, + ["customer_id"], + feature_table_df, + ["daily_transactions"], + feature_prefix="transactions__", + ) + + expected_joined_schema = StructType( + [ + StructField("customer_id", IntegerType()), + StructField("event_timestamp", TimestampType()), + StructField("transactions__daily_transactions", FloatType()), + ] + ) + expected_joined_data = [ + (1001, datetime(year=2020, month=8, day=31), None), + (1001, datetime(year=2020, month=9, day=1), 100.0,), + (1001, datetime(year=2020, month=9, day=2), 200.0,), + (1001, datetime(year=2020, month=9, day=3), 200.0,), + (2001, datetime(year=2020, month=9, day=2), 400.0,), + (3001, datetime(year=2020, month=9, day=1), None), + ] + expected_joined_df = spark.createDataFrame( + spark.sparkContext.parallelize(expected_joined_data), expected_joined_schema + ) + + assert_dataframe_equal(joined_df, expected_joined_df) + + +def test_join_with_max_age( + spark: SparkSession, + single_entity_schema: StructType, + customer_feature_schema: StructType, +): + entity_data = [ + (1001, datetime(year=2020, month=9, day=1)), + (1001, datetime(year=2020, month=9, day=3)), + (2001, datetime(year=2020, month=9, day=2)), + ] + entity_df = spark.createDataFrame( + spark.sparkContext.parallelize(entity_data), single_entity_schema + ) + + feature_table_data = [ + ( + 1001, + datetime(year=2020, month=9, day=1), + datetime(year=2020, month=9, day=1), + 100.0, + ), + ( + 2001, + datetime(year=2020, month=9, day=1), + datetime(year=2020, month=9, day=1), + 200.0, + ), + ] + feature_table_df = spark.createDataFrame( + spark.sparkContext.parallelize(feature_table_data), customer_feature_schema + ) + + joined_df = as_of_join( + entity_df, + ["customer_id"], + feature_table_df, + ["daily_transactions"], + feature_prefix="transactions__", + max_age=86400, + ) + + expected_joined_schema = StructType( + [ + StructField("customer_id", IntegerType()), + StructField("event_timestamp", TimestampType()), + StructField("transactions__daily_transactions", FloatType()), + ] + ) + expected_joined_data = [ + (1001, datetime(year=2020, month=9, day=1), 100.0,), + (1001, datetime(year=2020, month=9, day=3), None), + (2001, datetime(year=2020, month=9, day=2), 200.0,), + ] + expected_joined_df = spark.createDataFrame( + spark.sparkContext.parallelize(expected_joined_data), expected_joined_schema + ) + + assert_dataframe_equal(joined_df, expected_joined_df) + + +def test_join_with_composite_entity( + spark: SparkSession, + composite_entity_schema: StructType, + rating_feature_schema: StructType, +): + entity_data = [ + (1001, 8001, datetime(year=2020, month=9, day=1)), + (1001, 8002, datetime(year=2020, month=9, day=3)), + (1001, 8003, datetime(year=2020, month=9, day=1)), + (2001, 8001, datetime(year=2020, month=9, day=2)), + ] + entity_df = spark.createDataFrame( + spark.sparkContext.parallelize(entity_data), composite_entity_schema + ) + + feature_table_data = [ + ( + 1001, + 8001, + datetime(year=2020, month=9, day=1), + datetime(year=2020, month=9, day=1), + 3.0, + 5.0, + ), + ( + 1001, + 8002, + datetime(year=2020, month=9, day=1), + datetime(year=2020, month=9, day=1), + 4.0, + 3.0, + ), + ( + 2001, + 8001, + datetime(year=2020, month=9, day=1), + datetime(year=2020, month=9, day=1), + 4.0, + 4.5, + ), + ] + feature_table_df = spark.createDataFrame( + spark.sparkContext.parallelize(feature_table_data), rating_feature_schema, + ) + + joined_df = as_of_join( + entity_df, + ["customer_id", "driver_id"], + feature_table_df, + ["customer_rating", "driver_rating"], + feature_prefix="ratings__", + max_age=86400, + ) + + expected_joined_schema = StructType( + [ + StructField("customer_id", IntegerType()), + StructField("driver_id", IntegerType()), + StructField("event_timestamp", TimestampType()), + StructField("ratings__customer_rating", FloatType()), + StructField("ratings__driver_rating", FloatType()), + ] + ) + expected_joined_data = [ + (1001, 8001, datetime(year=2020, month=9, day=1), 3.0, 5.0,), + (1001, 8002, datetime(year=2020, month=9, day=3), None, None), + (1001, 8003, datetime(year=2020, month=9, day=1), None, None), + (2001, 8001, datetime(year=2020, month=9, day=2), 4.0, 4.5,), + ] + expected_joined_df = spark.createDataFrame( + spark.sparkContext.parallelize(expected_joined_data), expected_joined_schema + ) + + assert_dataframe_equal(joined_df, expected_joined_df) + + +def test_select_subset_of_columns_as_entity_primary_keys( + spark: SparkSession, + composite_entity_schema: StructType, + customer_feature_schema: StructType, +): + entity_data = [ + (1001, 8001, datetime(year=2020, month=9, day=2)), + (2001, 8002, datetime(year=2020, month=9, day=2)), + ] + entity_df = spark.createDataFrame( + spark.sparkContext.parallelize(entity_data), composite_entity_schema + ) + + feature_table_data = [ + ( + 1001, + datetime(year=2020, month=9, day=1), + datetime(year=2020, month=9, day=2), + 100.0, + ), + ( + 2001, + datetime(year=2020, month=9, day=1), + datetime(year=2020, month=9, day=1), + 400.0, + ), + ] + feature_table_df = spark.createDataFrame( + spark.sparkContext.parallelize(feature_table_data), customer_feature_schema + ) + + joined_df = as_of_join( + entity_df, + ["customer_id"], + feature_table_df, + ["daily_transactions"], + feature_prefix="transactions__", + ) + + expected_joined_schema = StructType( + [ + StructField("customer_id", IntegerType()), + StructField("driver_id", IntegerType()), + StructField("event_timestamp", TimestampType()), + StructField("transactions__daily_transactions", FloatType()), + ] + ) + expected_joined_data = [ + (1001, 8001, datetime(year=2020, month=9, day=2), 100.0,), + (2001, 8002, datetime(year=2020, month=9, day=2), 400.0,), + ] + expected_joined_df = spark.createDataFrame( + spark.sparkContext.parallelize(expected_joined_data), expected_joined_schema + ) + + assert_dataframe_equal(joined_df, expected_joined_df) + + +def test_multiple_join( + spark: SparkSession, + composite_entity_schema: StructType, + customer_feature_schema: StructType, + driver_feature_schema: StructType, +): + query_conf: List[Dict[str, Any]] = [ + { + "table": "transactions", + "features": ["daily_transactions"], + "join": ["customer_id"], + "max_age": 86400, + }, + { + "table": "bookings", + "features": ["completed_bookings"], + "join": ["driver_id"], + }, + ] + + entity_data = [ + (1001, 8001, datetime(year=2020, month=9, day=2)), + (1001, 8002, datetime(year=2020, month=9, day=2)), + (2001, 8002, datetime(year=2020, month=9, day=3)), + ] + entity_df = spark.createDataFrame( + spark.sparkContext.parallelize(entity_data), composite_entity_schema + ) + + customer_table_data = [ + ( + 1001, + datetime(year=2020, month=9, day=1), + datetime(year=2020, month=9, day=1), + 100.0, + ), + ( + 2001, + datetime(year=2020, month=9, day=1), + datetime(year=2020, month=9, day=1), + 200.0, + ), + ] + customer_table_df = spark.createDataFrame( + spark.sparkContext.parallelize(customer_table_data), customer_feature_schema + ) + + driver_table_data = [ + ( + 8001, + datetime(year=2020, month=8, day=31), + datetime(year=2020, month=8, day=31), + 200, + ), + ( + 8001, + datetime(year=2020, month=9, day=1), + datetime(year=2020, month=9, day=1), + 300, + ), + ( + 8002, + datetime(year=2020, month=9, day=1), + datetime(year=2020, month=9, day=1), + 600, + ), + ( + 8002, + datetime(year=2020, month=9, day=1), + datetime(year=2020, month=9, day=2), + 500, + ), + ] + driver_table_df = spark.createDataFrame( + spark.sparkContext.parallelize(driver_table_data), driver_feature_schema + ) + + tables = {"transactions": customer_table_df, "bookings": driver_table_df} + + joined_df = join_entity_to_feature_tables(query_conf, entity_df, tables) + + expected_joined_schema = StructType( + [ + StructField("customer_id", IntegerType()), + StructField("driver_id", IntegerType()), + StructField("event_timestamp", TimestampType()), + StructField("transactions__daily_transactions", FloatType()), + StructField("bookings__completed_bookings", IntegerType()), + ] + ) + + expected_joined_data = [ + (1001, 8001, datetime(year=2020, month=9, day=2), 100.0, 300,), + (1001, 8002, datetime(year=2020, month=9, day=2), 100.0, 500,), + (2001, 8002, datetime(year=2020, month=9, day=3), None, 500,), + ] + expected_joined_df = spark.createDataFrame( + spark.sparkContext.parallelize(expected_joined_data), expected_joined_schema + ) + + assert_dataframe_equal(joined_df, expected_joined_df) + + +def test_historical_feature_retrieval(spark): + test_data_dir = path.join(pathlib.Path(__file__).parent.absolute(), "data") + batch_retrieval_conf = { + "entity": { + "format": "csv", + "path": f"file://{path.join(test_data_dir, 'customer_driver_pairs.csv')}", + "options": {"inferSchema": "true", "header": "true"}, + "dtypes": {"customer_id": "int", "driver_id": "int"}, + }, + "tables": [ + { + "format": "csv", + "path": f"file://{path.join(test_data_dir, 'bookings.csv')}", + "name": "bookings", + "options": {"inferSchema": "true", "header": "true"}, + "dtypes": {"driver_id": "int"}, + }, + { + "format": "csv", + "path": f"file://{path.join(test_data_dir, 'transactions.csv')}", + "name": "transactions", + "options": {"inferSchema": "true", "header": "true"}, + "dtypes": {"customer_id": "int"}, + }, + ], + "queries": [ + { + "table": "transactions", + "features": ["daily_transactions"], + "join": ["customer_id"], + "max_age": 86400, + }, + { + "table": "bookings", + "features": ["completed_bookings"], + "join": ["driver_id"], + }, + ], + } + + joined_df = retrieve_historical_features(spark, batch_retrieval_conf) + + expected_joined_schema = StructType( + [ + StructField("customer_id", IntegerType()), + StructField("driver_id", IntegerType()), + StructField("event_timestamp", TimestampType()), + StructField("transactions__daily_transactions", FloatType()), + StructField("bookings__completed_bookings", IntegerType()), + ] + ) + + expected_joined_data = [ + (1001, 8001, datetime(year=2020, month=9, day=2), 100.0, 300,), + (1001, 8002, datetime(year=2020, month=9, day=2), 100.0, 500,), + (1001, 8002, datetime(year=2020, month=9, day=3), None, 500,), + (2001, 8002, datetime(year=2020, month=9, day=3), None, 500,), + (2001, 8002, datetime(year=2020, month=9, day=4), None, 500,), + ] + expected_joined_df = spark.createDataFrame( + spark.sparkContext.parallelize(expected_joined_data), expected_joined_schema + ) + + assert_dataframe_equal(joined_df, expected_joined_df) + + +def test_historical_feature_retrieval_with_mapping(spark): + test_data_dir = path.join(pathlib.Path(__file__).parent.absolute(), "data") + retrieval_conf = { + "entity": { + "format": "csv", + "path": f"file://{path.join(test_data_dir, 'column_mapping_test_entity.csv')}", + "options": {"inferSchema": "true", "header": "true"}, + "col_mapping": {"id": "customer_id"}, + "dtypes": {"customer_id": "int"}, + }, + "tables": [ + { + "format": "csv", + "path": f"file://{path.join(test_data_dir, 'column_mapping_test_feature.csv')}", + "name": "bookings", + "options": {"inferSchema": "true", "header": "true"}, + "col_mapping": { + "datetime": "event_timestamp", + "created_datetime": "created_timestamp", + }, + "dtypes": {"customer_id": "int"}, + }, + ], + "queries": [ + { + "table": "bookings", + "features": ["total_bookings"], + "join": ["customer_id"], + } + ], + } + + joined_df = retrieve_historical_features(spark, retrieval_conf) + + expected_joined_schema = StructType( + [ + StructField("customer_id", IntegerType()), + StructField("event_timestamp", TimestampType()), + StructField("bookings__total_bookings", IntegerType()), + ] + ) + + expected_joined_data = [ + (1001, datetime(year=2020, month=9, day=2), 200), + (1001, datetime(year=2020, month=9, day=3), 200), + (2001, datetime(year=2020, month=9, day=4), 600), + (2001, datetime(year=2020, month=9, day=4), 600), + (3001, datetime(year=2020, month=9, day=4), 700), + ] + expected_joined_df = spark.createDataFrame( + spark.sparkContext.parallelize(expected_joined_data), expected_joined_schema + ) + + assert_dataframe_equal(joined_df, expected_joined_df) + + +def test_large_historical_feature_retrieval( + spark, large_entity_csv_file, large_feature_csv_file +): + nr_rows = 1000 + start_datetime = datetime(year=2020, month=8, day=31) + expected_join_data = [ + (1000 + i, start_datetime + timedelta(days=i), i * 10) for i in range(nr_rows) + ] + expected_join_data_schema = StructType( + [ + StructField("customer_id", IntegerType()), + StructField("event_timestamp", TimestampType()), + StructField("feature__total_bookings", IntegerType()), + ] + ) + + expected_join_data_df = spark.createDataFrame( + spark.sparkContext.parallelize(expected_join_data), expected_join_data_schema + ) + + retrieval_conf = { + "entity": { + "format": "csv", + "path": f"file://{large_entity_csv_file}", + "options": {"inferSchema": "true", "header": "true"}, + }, + "tables": [ + { + "format": "csv", + "path": f"file://{large_feature_csv_file}", + "name": "feature", + "options": {"inferSchema": "true", "header": "true"}, + }, + ], + "queries": [ + { + "table": "feature", + "features": ["total_bookings"], + "join": ["customer_id"], + } + ], + } + + joined_df = retrieve_historical_features(spark, retrieval_conf) + assert_dataframe_equal(joined_df, expected_join_data_df) + + +def test_schema_verification(spark): + entity_schema = StructType( + [ + StructField("customer_id", IntegerType()), + StructField("event_timestamp", TimestampType()), + ] + ) + + entity_data = [ + (1001, datetime(year=2020, month=9, day=2)), + ] + + entity_df = spark.createDataFrame( + spark.sparkContext.parallelize(entity_data), entity_schema + ) + + with pytest.raises(MissingColumnError): + verify_schema(entity_df, {"driver_id": "int"}, False) + + with pytest.raises(SchemaMismatchError): + verify_schema(entity_df, {"customer_id": "string"}, False) + + with pytest.raises(MissingColumnError): + verify_schema(entity_df.drop("event_timestamp"), {"customer_id": "int"}, False) + + feature_schema = StructType( + [ + StructField("customer_id", IntegerType()), + StructField("event_timestamp", TimestampType()), + StructField("created_timestamp", TimestampType()), + ] + ) + + feature_data = [ + ( + 1001, + datetime(year=2020, month=9, day=2), + datetime(year=2020, month=9, day=2), + ), + ] + + feature_df = spark.createDataFrame( + spark.sparkContext.parallelize(feature_data), feature_schema + ) + + with pytest.raises(MissingColumnError): + verify_schema( + feature_df.drop("created_timestamp"), {"customer_id": "int"}, True + ) + + with pytest.raises(TimestampColumnError): + verify_schema( + feature_df.drop("created_timestamp").withColumn( + "created_timestamp", lit("test") + ), + {"customer_id": "int"}, + True, + )