def test_automapper_filter_and_transform(spark_session: SparkSession) -> None:
clean_spark_session(spark_session)
data_dir: Path = Path(__file__).parent.joinpath("./")
data_json_file: Path = data_dir.joinpath("data.json")
source_df: DataFrame = spark_session.read.json(str(data_json_file),
multiLine=True)
source_df.createOrReplaceTempView("patients")
source_df.show(truncate=False)
# Act
mapper = AutoMapper(view="members", source_view="patients").complex(
MyObject(age=A.transform(
A.filter(column=A.column("identifier"),
func=lambda x: x["use"] == lit("usual")),
A.complex(bar=A.field("value"), bar2=A.field("system")))))
assert isinstance(mapper, AutoMapper)
sql_expressions: Dict[str, Column] = mapper.get_column_specs(
source_df=source_df)
for column_name, sql_expression in sql_expressions.items():
print(f"{column_name}: {sql_expression}")
assert str(sql_expressions["age"]) == str(
transform(
filter("b.identifier", lambda x: x["use"] == lit("usual")),
lambda x: struct(x["value"].alias("bar"), x["system"].alias("bar2")
)).alias("age"))
result_df: DataFrame = mapper.transform(df=source_df)
result_df.show(truncate=False)
def get_column_spec(self, source_df: Optional[DataFrame],
current_column: Optional[Column]) -> Column:
if isinstance(
self.value, str
): # if the src column is just string then consider it a sql expression
return array(lit(self.value))
if isinstance(self.value,
list): # if the src column is a list then iterate
return filter(array(
[
self.get_value(item, source_df=source_df, current_column=current_column)
for item in self.value
]
), lambda x: x.isNotNull()) \
if self.remove_nulls \
else array(
[
self.get_value(item, source_df=source_df, current_column=current_column)
for item in self.value
]
)
# if value is an AutoMapper then ask it for its column spec
if isinstance(self.value, AutoMapperDataTypeBase):
child: AutoMapperDataTypeBase = self.value
return child.get_column_spec(source_df=source_df,
current_column=current_column)
raise ValueError(f"value: {self.value} is neither str nor AutoMapper")
def test_auto_mapper_array_multiple_items_with_null(
spark_session: SparkSession) -> None:
# Arrange
spark_session.createDataFrame([
(1, 'Qureshi', 'Imran'),
(2, 'Vidal', 'Michael'),
], ['member_id', 'last_name', 'first_name'
]).createOrReplaceTempView("patients")
source_df: DataFrame = spark_session.table("patients")
df: DataFrame = source_df.select("member_id")
df.createOrReplaceTempView("members")
# Act
mapper = AutoMapper(
view="members",
source_view="patients",
keys=["member_id"],
drop_key_columns=False).columns(
dst2=AutoMapperList(["address1", "address2", None]))
assert isinstance(mapper, AutoMapper)
sql_expressions: Dict[str, Column] = mapper.get_column_specs(
source_df=source_df)
for column_name, sql_expression in sql_expressions.items():
print(f"{column_name}: {sql_expression}")
assert str(sql_expressions["dst2"]) == str(
filter(array(lit("address1"), lit("address2"), lit(None)),
lambda x: x.isNotNull()).alias("dst2"))
result_df: DataFrame = mapper.transform(df=df)
# Assert
result_df.printSchema()
result_df.show()
assert result_df.where("member_id == 1").select(
"dst2").collect()[0][0][0] == "address1"
assert result_df.where("member_id == 1").select(
"dst2").collect()[0][0][1] == "address2"
assert result_df.where("member_id == 2").select(
"dst2").collect()[0][0][0] == "address1"
assert result_df.where("member_id == 2").select(
"dst2").collect()[0][0][1] == "address2"
def test_automapper_select_one(spark_session: SparkSession) -> None:
clean_spark_session(spark_session)
data_dir: Path = Path(__file__).parent.joinpath("./")
data_json_file: Path = data_dir.joinpath("data.json")
source_df: DataFrame = spark_session.read.json(
str(data_json_file), multiLine=True
)
source_df.createOrReplaceTempView("patients")
source_df.show(truncate=False)
# Act
mapper = AutoMapper(view="members", source_view="patients").columns(
age=A.column("identifier").filter(
lambda x: x["system"] == "http://hl7.org/fhir/sid/us-npi"
).select_one(A.field("_.value"))
)
assert isinstance(mapper, AutoMapper)
sql_expressions: Dict[str, Column] = mapper.get_column_specs(
source_df=source_df
)
for column_name, sql_expression in sql_expressions.items():
print(f"{column_name}: {sql_expression}")
assert str(sql_expressions["age"]) == str(
transform(
filter(
"b.identifier",
lambda x: x["system"] == lit("http://hl7.org/fhir/sid/us-npi")
), lambda x: x["value"]
)[0].alias("age")
)
result_df: DataFrame = mapper.transform(df=source_df)
result_df.show(truncate=False)
assert result_df.select("age").collect()[0][0] == "1730325416"
assert result_df.select("age").collect()[1][0] == "1467734301"
def get_column_spec(self, source_df: Optional[DataFrame],
current_column: Optional[Column]) -> Column:
return filter(
self.column.get_column_spec(source_df=source_df,
current_column=current_column),
self.func)
def test_auto_mapper_columns(spark_session: SparkSession) -> None:
# Arrange
spark_session.createDataFrame([
(1, 'Qureshi', 'Imran'),
(2, 'Vidal', 'Michael'),
], ['member_id', 'last_name', 'first_name'
]).createOrReplaceTempView("patients")
source_df: DataFrame = spark_session.table("patients")
df = source_df.select("member_id")
df.createOrReplaceTempView("members")
# Act
mapper = AutoMapper(view="members",
source_view="patients",
keys=["member_id"],
drop_key_columns=False).columns(
dst1="src1",
dst2=AutoMapperList(["address1"]),
dst3=AutoMapperList(["address1", "address2"]),
dst4=AutoMapperList([
A.complex(use="usual",
family=A.column("last_name"))
]))
assert isinstance(mapper, AutoMapper)
sql_expressions: Dict[str, Column] = mapper.get_column_specs(
source_df=source_df)
for column_name, sql_expression in sql_expressions.items():
print(f"{column_name}: {sql_expression}")
# Assert
assert len(sql_expressions) == 4
assert str(sql_expressions["dst1"]) == str(lit("src1").alias("dst1"))
assert str(sql_expressions["dst2"]) == str(
filter(array(lit("address1")), lambda x: x.isNotNull()).alias("dst2"))
assert str(sql_expressions["dst3"]) == str(
filter(array(lit("address1"), lit("address2")),
lambda x: x.isNotNull()).alias("dst3"))
assert str(sql_expressions["dst4"]) == str(
filter(
array(
struct(
lit("usual").alias("use"),
col("b.last_name").alias("family"))),
lambda x: x.isNotNull()).alias("dst4"))
result_df: DataFrame = mapper.transform(df=df)
# Assert
result_df.printSchema()
result_df.show()
assert len(result_df.columns) == 5
assert result_df.where("member_id == 1").select(
"dst1").collect()[0][0] == "src1"
assert result_df.where("member_id == 1").select(
"dst2").collect()[0][0][0] == "address1"
assert result_df.where("member_id == 1").select(
"dst3").collect()[0][0][0] == "address1"
assert result_df.where("member_id == 1").select(
"dst3").collect()[0][0][1] == "address2"
assert result_df.where("member_id == 1").select(
"dst4").collect()[0][0][0][0] == "usual"
assert result_df.where("member_id == 1").select(
"dst4").collect()[0][0][0][1] == "Qureshi"