update tests

Author: Davies Liu

python/pyspark/sql/context.py

@@ -557,6 +557,7 @@ def read(self):
         ::note: Experimental
 
         >>> sqlContext.read
+        <pyspark.sql.readwriter.DataFrameReader object at ...>
         """
         return DataFrameReader(self)
 

python/pyspark/sql/dataframe.py

@@ -193,6 +193,11 @@ def write(self):
         into external storage.
 
         :return :class:`DataFrameWriter`
+
+        ::note: Experimental
+
+        >>> df.write
+        <pyspark.sql.readwriter.DataFrameWriter object at ...>
         """
         return DataFrameWriter(self)
 

python/pyspark/sql/readwriter.py

@@ -54,7 +54,7 @@ def load(self, path=None, format=None, schema=None, **options):
         if schema is not None:
             if not isinstance(schema, StructType):
                 raise TypeError("schema should be StructType")
-            jschema = self.sqlContext._ssql_ctx.parseDataType(schema.json())
+            jschema = self._sqlContext._ssql_ctx.parseDataType(schema.json())
             jreader = jreader.schema(jschema)
         for k in options:
             jreader = jreader.option(k, options[k])
@@ -79,7 +79,7 @@ def json(self, path, schema=None):
         >>> shutil.rmtree(jsonFile)
         >>> with open(jsonFile, 'w') as f:
         ...     f.writelines(jsonStrings)
-        >>> df1 = sqlContext.jsonFile(jsonFile)
+        >>> df1 = sqlContext.read.json(jsonFile)
         >>> df1.printSchema()
         root
          |-- field1: long (nullable = true)
@@ -92,7 +92,7 @@ def json(self, path, schema=None):
         ...     StructField("field2", StringType()),
         ...     StructField("field3",
         ...         StructType([StructField("field5", ArrayType(IntegerType()))]))])
-        >>> df2 = sqlContext.jsonFile(jsonFile, schema)
+        >>> df2 = sqlContext.read.json(jsonFile, schema)
         >>> df2.printSchema()
         root
          |-- field2: string (nullable = true)
@@ -103,7 +103,7 @@ def json(self, path, schema=None):
         if schema is None:
             jdf = self._jreader.json(path)
         else:
-            jschema = self.sqlContext._ssql_ctx.parseDataType(schema.json())
+            jschema = self._sqlContext._ssql_ctx.parseDataType(schema.json())
             jdf = self._jreader.schema(jschema).json(path)
         return self._df(jdf)
 
@@ -181,7 +181,8 @@ class DataFrameWriter(object):
     """
     def __init__(self, df):
         self._df = df
-        self._jwrite = df._df.write()
+        self._sqlContext = df.sql_ctx
+        self._jwrite = df._jdf.write()
 
     def save(self, path=None, format=None, mode="error", **options):
         """

python/pyspark/sql/tests.py

@@ -480,29 +480,29 @@ def test_save_and_load(self):
         df = self.df
         tmpPath = tempfile.mkdtemp()
         shutil.rmtree(tmpPath)
-        df.save(tmpPath, "org.apache.spark.sql.json", "error")
-        actual = self.sqlCtx.load(tmpPath, "org.apache.spark.sql.json")
-        self.assertTrue(sorted(df.collect()) == sorted(actual.collect()))
+        df.write.json(tmpPath)
+        actual = self.sqlCtx.read.json(tmpPath)
+        self.assertEqual(sorted(df.collect()), sorted(actual.collect()))
 
         schema = StructType([StructField("value", StringType(), True)])
-        actual = self.sqlCtx.load(tmpPath, "org.apache.spark.sql.json", schema)
-        self.assertTrue(sorted(df.select("value").collect()) == sorted(actual.collect()))
+        actual = self.sqlCtx.read.json(tmpPath, schema)
+        self.assertEqual(sorted(df.select("value").collect()), sorted(actual.collect()))
 
-        df.save(tmpPath, "org.apache.spark.sql.json", "overwrite")
-        actual = self.sqlCtx.load(tmpPath, "org.apache.spark.sql.json")
-        self.assertTrue(sorted(df.collect()) == sorted(actual.collect()))
+        df.write.json(tmpPath, "overwrite")
+        actual = self.sqlCtx.read.json(tmpPath)
+        self.assertEqual(sorted(df.collect()), sorted(actual.collect()))
 
-        df.save(source="org.apache.spark.sql.json", mode="overwrite", path=tmpPath,
-                noUse="this options will not be used in save.")
-        actual = self.sqlCtx.load(source="org.apache.spark.sql.json", path=tmpPath,
-                                  noUse="this options will not be used in load.")
-        self.assertTrue(sorted(df.collect()) == sorted(actual.collect()))
+        df.write.save(format="json", mode="overwrite", path=tmpPath,
+                      noUse="this options will not be used in save.")
+        actual = self.sqlCtx.read.load(format="json", path=tmpPath,
+                                       noUse="this options will not be used in load.")
+        self.assertEqual(sorted(df.collect()), sorted(actual.collect()))
 
         defaultDataSourceName = self.sqlCtx.getConf("spark.sql.sources.default",
                                                     "org.apache.spark.sql.parquet")
         self.sqlCtx.sql("SET spark.sql.sources.default=org.apache.spark.sql.json")
         actual = self.sqlCtx.load(path=tmpPath)
-        self.assertTrue(sorted(df.collect()) == sorted(actual.collect()))
+        self.assertEqual(sorted(df.collect()), sorted(actual.collect()))
         self.sqlCtx.sql("SET spark.sql.sources.default=" + defaultDataSourceName)
 
         shutil.rmtree(tmpPath)