Compare rows of two dataframes to find the matching column count of 1's

Note : As mentioned by pault, this will work better if you have unique row indices that connect both dataframes. Otherwise, the row orders may not be guaranteed in some Spark operations.

(1) Setup the environment and some sample data.

import numpy as np

from pyspark.ml.feature import VectorAssembler
from pyspark.sql import functions as F

df1 = spark.createDataFrame([
    (0, 0, 1),
    (1, 0, 0),
    (0, 0, 1)
], ["column1", "column2", "column3"])

df2 = spark.createDataFrame([
    (0, 0, 0),
    (1, 0, 1),
    (0, 0, 1)
], ["column1", "column2", "column3"])

(2) Collect all columns into a Spark vector.

assembler = VectorAssembler(
    inputCols=["column1", "column2", "column3"],
    outputCol="merged_col")

df1_merged = assembler.transform(df1)
df2_merged = assembler.transform(df2)

df1_merged.show()

+-------+-------+-------+-------------+
|column1|column2|column3|   merged_col|
+-------+-------+-------+-------------+
|      0|      0|      1|[0.0,0.0,1.0]|
|      1|      0|      0|[1.0,0.0,0.0]|
|      0|      0|      1|[0.0,0.0,1.0]|
+-------+-------+-------+-------------+

(3) Get the row and column index of non-zero elements. Using numpy.nonzero() on RDD of Spark Vector.

def get_nonzero_index(args):
    (row, index) = args
    np_arr = np.array(row.merged_col)
    return (index, np_arr.nonzero()[0].tolist())

df1_ind_rdd = df1_merged.rdd.zipWithIndex().map(get_nonzero_index)
df2_ind_rdd = df2_merged.rdd.zipWithIndex().map(get_nonzero_index)

df1_ind_rdd.collect()
[(0, [2]), (1, [0]), (2, [2])]

df2_ind_rdd.collect()
[(0, []), (1, [0, 2]), (2, [2])]

(4) You can then do your comparison on these 2 Python lists easily.

Note that this method will not be efficient (due to collect) if the number of rows you have is very large. In that case, you will want to do all processing in Spark by doing a join on the 2 dataframes.

(5) To do the matching purely in Spark, you can try the methods below that rely on a join on row index.

df1_index = spark.createDataFrame(df1_ind_rdd, ["row_index_1", "column_index_1"])
df2_index = spark.createDataFrame(df2_ind_rdd, ["row_index_2", "column_index_2"])

df_joined = df1_index.join(df2_index, df1_index.row_index_1 == df2_index.row_index_2)

Then expand the list so that we get an element on each row.

df_exploded = df_joined.withColumn("column_index_exp_1", F.explode(df_joined.column_index_1))\
                            .withColumn("column_index_exp_2", F.explode(df_joined.column_index_2))

Check for match between the two columns and finally convert into integer for summing.

df_match_bool = df_exploded.withColumn("match_bool", df_exploded.column_index_exp_1 == df_exploded.column_index_exp_2)

df_match_int = df_match_bool.withColumn("match_integer", df_match_bool.match_bool.cast("long"))

df_match_bool.show()
+-----------+--------------+-----------+--------------+------------------+------------------+----------+
|row_index_1|column_index_1|row_index_2|column_index_2|column_index_exp_1|column_index_exp_2|match_bool|
+-----------+--------------+-----------+--------------+------------------+------------------+----------+
|          1|           [0]|          1|        [0, 2]|                 0|                 0|      true|
|          1|           [0]|          1|        [0, 2]|                 0|                 2|     false|
|          2|           [2]|          2|           [2]|                 2|                 2|      true|
+-----------+--------------+-----------+--------------+------------------+------------------+----------+

df_match_int.groupBy().sum("match_integer").collect()[0][0]
2