From 02bbe73118a39e2fb378aa2002449367a92f6d67 Mon Sep 17 00:00:00 2001
From: zero323 <zero323@users.noreply.github.com>
Date: Wed, 3 May 2017 19:15:28 -0700
Subject: [PATCH] [SPARK-20584][PYSPARK][SQL] Python generic hint support
## What changes were proposed in this pull request?
Adds `hint` method to PySpark `DataFrame`.
## How was this patch tested?
Unit tests, doctests.
Author: zero323 <zero323@users.noreply.github.com>
Closes #17850 from zero323/SPARK-20584.
---
python/pyspark/sql/dataframe.py | 29 +++++++++++++++++++++++++++++
python/pyspark/sql/tests.py | 16 ++++++++++++++++
2 files changed, 45 insertions(+)
diff --git a/python/pyspark/sql/dataframe.py b/python/pyspark/sql/dataframe.py
index ab6d35bfa7..7b67985f2b 100644
--- a/python/pyspark/sql/dataframe.py
+++ b/python/pyspark/sql/dataframe.py
@@ -380,6 +380,35 @@ class DataFrame(object):
jdf = self._jdf.withWatermark(eventTime, delayThreshold)
return DataFrame(jdf, self.sql_ctx)
+ @since(2.2)
+ def hint(self, name, *parameters):
+ """Specifies some hint on the current DataFrame.
+
+ :param name: A name of the hint.
+ :param parameters: Optional parameters.
+ :return: :class:`DataFrame`
+
+ >>> df.join(df2.hint("broadcast"), "name").show()
+ +----+---+------+
+ |name|age|height|
+ +----+---+------+
+ | Bob| 5| 85|
+ +----+---+------+
+ """
+ if len(parameters) == 1 and isinstance(parameters[0], list):
+ parameters = parameters[0]
+
+ if not isinstance(name, str):
+ raise TypeError("name should be provided as str, got {0}".format(type(name)))
+
+ for p in parameters:
+ if not isinstance(p, str):
+ raise TypeError(
+ "all parameters should be str, got {0} of type {1}".format(p, type(p)))
+
+ jdf = self._jdf.hint(name, self._jseq(parameters))
+ return DataFrame(jdf, self.sql_ctx)
+
@since(1.3)
def count(self):
"""Returns the number of rows in this :class:`DataFrame`.
diff --git a/python/pyspark/sql/tests.py b/python/pyspark/sql/tests.py
index ce4abf8fb7..f644624f7f 100644
--- a/python/pyspark/sql/tests.py
+++ b/python/pyspark/sql/tests.py
@@ -1906,6 +1906,22 @@ class SQLTests(ReusedPySparkTestCase):
# planner should not crash without a join
broadcast(df1)._jdf.queryExecution().executedPlan()
+ def test_generic_hints(self):
+ from pyspark.sql import DataFrame
+
+ df1 = self.spark.range(10e10).toDF("id")
+ df2 = self.spark.range(10e10).toDF("id")
+
+ self.assertIsInstance(df1.hint("broadcast"), DataFrame)
+ self.assertIsInstance(df1.hint("broadcast", []), DataFrame)
+
+ # Dummy rules
+ self.assertIsInstance(df1.hint("broadcast", "foo", "bar"), DataFrame)
+ self.assertIsInstance(df1.hint("broadcast", ["foo", "bar"]), DataFrame)
+
+ plan = df1.join(df2.hint("broadcast"), "id")._jdf.queryExecution().executedPlan()
+ self.assertEqual(1, plan.toString().count("BroadcastHashJoin"))
+
def test_toDF_with_schema_string(self):
data = [Row(key=i, value=str(i)) for i in range(100)]
rdd = self.sc.parallelize(data, 5)
--
GitLab