def spark_mock_customer(pd_mock_customer): ps = pytest.importorskip("pyspark.pandas", reason="Spark not installed, skipping") dataframes = {} for df in pd_mock_customer.dataframes: cleaned_df = pd_to_spark_clean(df).reset_index(drop=True) dataframes[df.ww.name] = ( ps.from_pandas(cleaned_df), df.ww.index, df.ww.time_index, df.ww.logical_types, ) relationships = [ ( rel._parent_dataframe_name, rel._parent_column_name, rel._child_dataframe_name, rel._child_column_name, ) for rel in pd_mock_customer.relationships ] return ft.EntitySet( id=pd_mock_customer.id, dataframes=dataframes, relationships=relationships )
def spark_latlong_df(pd_latlong_df): ps = pytest.importorskip("pyspark.pandas", reason="Spark not installed, skipping") cleaned_df = pd_to_spark_clean(pd_latlong_df) pdf = ps.from_pandas(cleaned_df) return pdf
def test_add_dataframe_from_spark_df(pd_es): cleaned_df = pd_to_spark_clean(pd_es["log"]) log_spark = ps.from_pandas(cleaned_df) spark_es = EntitySet(id="spark_es") spark_es = spark_es.add_dataframe( dataframe_name="log_spark", dataframe=log_spark, index="id", time_index="datetime", logical_types=pd_es["log"].ww.logical_types, semantic_tags=get_df_tags(pd_es["log"]), ) pd.testing.assert_frame_equal( cleaned_df, spark_es["log_spark"].to_pandas(), check_like=True )
def spark_int_es(pd_int_es): ps = pytest.importorskip("pyspark.pandas", reason="Spark not installed, skipping") es = ft.EntitySet(id=pd_int_es.id) for df in pd_int_es.dataframes: cleaned_df = pd_to_spark_clean(df).reset_index(drop=True) spark_df = ps.from_pandas(cleaned_df) spark_df.ww.init(schema=df.ww.schema) es.add_dataframe(spark_df) for rel in pd_int_es.relationships: es.add_relationship( rel._parent_dataframe_name, rel._parent_column_name, rel._child_dataframe_name, rel._child_column_name, ) return es
def spark_home_games_es(pd_home_games_es): ps = pytest.importorskip("pyspark.pandas", reason="Spark not installed, skipping") dataframes = {} for df in pd_home_games_es.dataframes: spark_df = ps.from_pandas(pd_to_spark_clean(df)) spark_df.ww.init(schema=df.ww.schema) dataframes[df.ww.name] = (spark_df,) relationships = [ ( rel._parent_dataframe_name, rel._parent_column_name, rel._child_dataframe_name, rel._child_column_name, ) for rel in pd_home_games_es.relationships ] return ft.EntitySet( id=pd_home_games_es.id, dataframes=dataframes, relationships=relationships )