adding keys() and values()

databricks · May 17, 2024 · 9e53506 · 9e53506
1 parent 103913e
commit 9e53506
Show file tree

Hide file tree

Showing 4 changed files with 28 additions and 2 deletions.
diff --git a/README.md b/README.md
@@ -106,7 +106,7 @@ that open a pull request and we will review it.
 | isEmpty                           | :x:                |                                                                   |
 | isLocallyCheckpointed             | :x:                |                                                                   |
 | join                              | :x:                |                                                                   |
-| keys                              | :x:                |                                                                   |
+| keys                              | :white_check_mark: |                                                                   |
 | leftOuterJoin                     | :x:                |                                                                   |
 | localCheckpoint                   | :x:                |                                                                   |
 | lookup                            | :x:                |                                                                   |
@@ -159,7 +159,7 @@ that open a pull request and we will review it.
 | treeReduce                        | :x:                |                                                                   |
 | union                             | :x:                |                                                                   |
 | unpersist                         | :x:                |                                                                   |
-| values                            | :x:                |                                                                   |
+| values                            | :white_check_mark: |                                                                   |
 | variance                          | :x:                |                                                                   |
 | withResources                     | :x:                |                                                                   |
 | zip                               | :x:                |                                                                   |

diff --git a/congruity/rdd_adapter.py b/congruity/rdd_adapter.py
@@ -274,6 +274,16 @@ def func(iterator: Iterable[T]) -> Iterable[T]:
 
     fold.__doc__ = RDD.fold.__doc__
 
+    def keys(self) -> "RDDAdapter":
+        return self.map(lambda x: x[0])
+
+    keys.__doc__ = RDD.keys.__doc__
+
+    def values(self) -> "RDDAdapter":
+        return self.map(lambda x: x[1])
+
+    values.__doc__ = RDD.values.__doc__
+
     class WrappedIterator(Iterable):
         """This is a helper class that wraps the iterator of RecordBatches as returned by
         mapInArrow and converts it into an iterator of the underlaying values."""

diff --git a/tests/test_public_interface.py b/tests/test_public_interface.py
@@ -34,3 +34,7 @@ def test_spark_context_parallelize(spark_session: "SparkSession"):
         .collect()
     )
     assert result == [20000, 100000, 3000]
+
+    val = spark_session.sparkContext.parallelize(range(0, 5))
+    assert val.count() == 5
+    assert val.collect() == [0, 1, 2, 3, 4]
diff --git a/tests/test_rdd_adapter.py b/tests/test_rdd_adapter.py
@@ -130,3 +130,15 @@ def test_rdd_sum(spark_session: "SparkSession"):
 
     vals = df.rdd.map(lambda x: x[0]).sum()
     assert vals == 45
+
+
+def test_rdd_keys(spark_session: "SparkSession"):
+    monkey_patch_spark()
+    rdd = spark_session.sparkContext.parallelize([(1, 2), (3, 4)]).keys()
+    assert rdd.collect() == [1, 3]
+
+
+def test_rdd_values(spark_session: "SparkSession"):
+    monkey_patch_spark()
+    rdd = spark_session.sparkContext.parallelize([(1, 2), (3, 4)]).values()
+    assert rdd.collect() == [2, 4]