aboutsummaryrefslogtreecommitdiff
path: root/python/pyspark/sql/dataframe.py
diff options
context:
space:
mode:
Diffstat (limited to 'python/pyspark/sql/dataframe.py')
-rw-r--r--python/pyspark/sql/dataframe.py21
1 files changed, 18 insertions, 3 deletions
diff --git a/python/pyspark/sql/dataframe.py b/python/pyspark/sql/dataframe.py
index 7673153abe..03b01a1136 100644
--- a/python/pyspark/sql/dataframe.py
+++ b/python/pyspark/sql/dataframe.py
@@ -1189,15 +1189,30 @@ class DataFrame(object):
@since(1.4)
@ignore_unicode_prefix
- def drop(self, colName):
+ def drop(self, col):
"""Returns a new :class:`DataFrame` that drops the specified column.
- :param colName: string, name of the column to drop.
+ :param col: a string name of the column to drop, or a
+ :class:`Column` to drop.
>>> df.drop('age').collect()
[Row(name=u'Alice'), Row(name=u'Bob')]
+
+ >>> df.drop(df.age).collect()
+ [Row(name=u'Alice'), Row(name=u'Bob')]
+
+ >>> df.join(df2, df.name == df2.name, 'inner').drop(df.name).collect()
+ [Row(age=5, height=85, name=u'Bob')]
+
+ >>> df.join(df2, df.name == df2.name, 'inner').drop(df2.name).collect()
+ [Row(age=5, name=u'Bob', height=85)]
"""
- jdf = self._jdf.drop(colName)
+ if isinstance(col, basestring):
+ jdf = self._jdf.drop(col)
+ elif isinstance(col, Column):
+ jdf = self._jdf.drop(col._jc)
+ else:
+ raise TypeError("col should be a string or a Column")
return DataFrame(jdf, self.sql_ctx)
@since(1.3)