From 1150a19b188a075166899fdb1e107b2ba1e505d8 Mon Sep 17 00:00:00 2001
From: Wenchen Fan <cloud0fan@outlook.com>
Date: Fri, 14 Aug 2015 14:09:46 -0700
Subject: [PATCH] [SPARK-8670] [SQL] Nested columns can't be referenced in
pyspark
This bug is caused by a wrong column-exist-check in `__getitem__` of pyspark dataframe. `DataFrame.apply` accepts not only top level column names, but also nested column name like `a.b`, so we should remove that check from `__getitem__`.
Author: Wenchen Fan <cloud0fan@outlook.com>
Closes #8202 from cloud-fan/nested.
---
python/pyspark/sql/dataframe.py | 2 --
python/pyspark/sql/tests.py | 4 +++-
sql/core/src/main/scala/org/apache/spark/sql/DataFrame.scala | 2 ++
3 files changed, 5 insertions(+), 3 deletions(-)
diff --git a/python/pyspark/sql/dataframe.py b/python/pyspark/sql/dataframe.py
index 09647ff6d0..da742d7ce7 100644
--- a/python/pyspark/sql/dataframe.py
+++ b/python/pyspark/sql/dataframe.py
@@ -722,8 +722,6 @@ class DataFrame(object):
[Row(age=5, name=u'Bob')]
"""
if isinstance(item, basestring):
- if item not in self.columns:
- raise IndexError("no such column: %s" % item)
jc = self._jdf.apply(item)
return Column(jc)
elif isinstance(item, Column):
diff --git a/python/pyspark/sql/tests.py b/python/pyspark/sql/tests.py
index 9b748101b5..13cf647b66 100644
--- a/python/pyspark/sql/tests.py
+++ b/python/pyspark/sql/tests.py
@@ -770,7 +770,7 @@ class SQLTests(ReusedPySparkTestCase):
self.assertTrue(isinstance(df['key'], Column))
self.assertTrue(isinstance(df[0], Column))
self.assertRaises(IndexError, lambda: df[2])
- self.assertRaises(IndexError, lambda: df["bad_key"])
+ self.assertRaises(AnalysisException, lambda: df["bad_key"])
self.assertRaises(TypeError, lambda: df[{}])
def test_column_name_with_non_ascii(self):
@@ -794,7 +794,9 @@ class SQLTests(ReusedPySparkTestCase):
df = self.sc.parallelize([Row(l=[1], r=Row(a=1, b="b"), d={"k": "v"})]).toDF()
self.assertEqual(1, df.select(df.l[0]).first()[0])
self.assertEqual(1, df.select(df.r["a"]).first()[0])
+ self.assertEqual(1, df.select(df["r.a"]).first()[0])
self.assertEqual("b", df.select(df.r["b"]).first()[0])
+ self.assertEqual("b", df.select(df["r.b"]).first()[0])
self.assertEqual("v", df.select(df.d["k"]).first()[0])
def test_infer_long_type(self):
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/DataFrame.scala b/sql/core/src/main/scala/org/apache/spark/sql/DataFrame.scala
index cf75e64e88..fd0ead4401 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/DataFrame.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/DataFrame.scala
@@ -634,6 +634,7 @@ class DataFrame private[sql](
/**
* Selects column based on the column name and return it as a [[Column]].
+ * Note that the column name can also reference to a nested column like `a.b`.
* @group dfops
* @since 1.3.0
*/
@@ -641,6 +642,7 @@ class DataFrame private[sql](
/**
* Selects column based on the column name and return it as a [[Column]].
+ * Note that the column name can also reference to a nested column like `a.b`.
* @group dfops
* @since 1.3.0
*/
--
GitLab