aboutsummaryrefslogtreecommitdiff
path: root/python
diff options
context:
space:
mode:
Diffstat (limited to 'python')
-rw-r--r--python/pyspark/sql/functions.py9
-rw-r--r--python/pyspark/sql/tests.py18
2 files changed, 27 insertions, 0 deletions
diff --git a/python/pyspark/sql/functions.py b/python/pyspark/sql/functions.py
index 26b8662718..fa04f4cd83 100644
--- a/python/pyspark/sql/functions.py
+++ b/python/pyspark/sql/functions.py
@@ -29,6 +29,7 @@ from pyspark.rdd import _prepare_for_python_RDD, ignore_unicode_prefix
from pyspark.serializers import PickleSerializer, AutoBatchedSerializer
from pyspark.sql.types import StringType
from pyspark.sql.column import Column, _to_java_column, _to_seq
+from pyspark.sql.dataframe import DataFrame
def _create_function(name, doc=""):
@@ -189,6 +190,14 @@ def approxCountDistinct(col, rsd=None):
return Column(jc)
+@since(1.6)
+def broadcast(df):
+ """Marks a DataFrame as small enough for use in broadcast joins."""
+
+ sc = SparkContext._active_spark_context
+ return DataFrame(sc._jvm.functions.broadcast(df._jdf), df.sql_ctx)
+
+
@since(1.4)
def coalesce(*cols):
"""Returns the first column that is not null.
diff --git a/python/pyspark/sql/tests.py b/python/pyspark/sql/tests.py
index 3e680f1030..645133b2b2 100644
--- a/python/pyspark/sql/tests.py
+++ b/python/pyspark/sql/tests.py
@@ -1075,6 +1075,24 @@ class SQLTests(ReusedPySparkTestCase):
self.assertRaises(TypeError, foo)
+ # add test for SPARK-10577 (test broadcast join hint)
+ def test_functions_broadcast(self):
+ from pyspark.sql.functions import broadcast
+
+ df1 = self.sqlCtx.createDataFrame([(1, "1"), (2, "2")], ("key", "value"))
+ df2 = self.sqlCtx.createDataFrame([(1, "1"), (2, "2")], ("key", "value"))
+
+ # equijoin - should be converted into broadcast join
+ plan1 = df1.join(broadcast(df2), "key")._jdf.queryExecution().executedPlan()
+ self.assertEqual(1, plan1.toString().count("BroadcastHashJoin"))
+
+ # no join key -- should not be a broadcast join
+ plan2 = df1.join(broadcast(df2))._jdf.queryExecution().executedPlan()
+ self.assertEqual(0, plan2.toString().count("BroadcastHashJoin"))
+
+ # planner should not crash without a join
+ broadcast(df1)._jdf.queryExecution().executedPlan()
+
class HiveContextSQLTests(ReusedPySparkTestCase):