Btukfyl

I have the following Dataframe

+--------------+---------------+   

|       SrcAddr|        DstAddr|   

+--------------+---------------+  

| 192.168.100.5| 192.168.220.16|  

| 192.168.100.5| 192.168.220.15|  

|192.168.220.15|  192.168.100.5|  

|192.168.220.16|  192.168.100.5|  

| 192.168.100.5| 192.168.220.15|  

|192.168.220.16|  192.168.100.5|  

| 192.168.220.9|  192.168.100.5|  

| 192.168.100.5|  192.168.220.9|  

| 192.168.220.9|  192.168.100.5|  

+--------------+---------------+  

containing source and destination address IPs.
I want to transform them in numerical index by means of StringIndexer, but I want to learn a common mapping between the columns.

Unfortunately StringIndexer does not provide such a rich interface in PySpark. Thus I found a workaround, but I wanted to know if there is a better way to do it.

What I have done is the following:

First, I compute the union between the two columns

src_addr_df = df.select(["SrcAddr"]).withColumnRenamed("SrcAddr", "Addr")  

dst_addr_df = df.select(["DstAddr"]).withColumnRenamed("DstAddr", "Addr")  

all_addr_df = src_addr_df.union(dst_addr_df)  

Then, I learned a common StringIndexer over the newly created DataFrame:

addrIndexer = StringIndexer(inputCol="Addr", outputCol="AddrIdx")  

addrModel = addrIndexer.fit(all_addr_df)  

Finally, I used the learned model to transform the original dataframe. This, is the tricky part because I need to rename the columns quite often to obtain the desired results:

df = addrModel.transform(df.withColumnRenamed("SrcAddr", "Addr")).withColumnRenamed("Addr", "SrcAddr").withColumnRenamed("AddrIdx", "SrcAddrIdx")



df = addrModel.transform(df.withColumnRenamed("DstAddr", "Addr")).withColumnRenamed("Addr", "DstAddr").withColumnRenamed("AddrIdx", "DstAddrIdx")

Thus, I'm wandering if there is the possibility to rather change the InputCol value of the StringIndexer, which would create a much readable code

Best regards,
Sandro