more working

Author: ahirreddy

python/pyspark/context.py

@@ -32,7 +32,7 @@
         PairDeserializer
 from pyspark.storagelevel import StorageLevel
 from pyspark import rdd
-from pyspark.rdd import RDD
+from pyspark.rdd import RDD, SchemaRDD
 
 from py4j.java_collections import ListConverter
 
@@ -462,6 +462,23 @@ def sparkUser(self):
         """
         return self._jsc.sc().sparkUser()
 
+class SQLContext:
+
+    def __init__(self, sparkContext):
+        self._sc = sparkContext
+        self._jsc = self._sc._jsc
+        self._jvm = self._sc._jvm
+        self._jsql_ctx = self._jvm.JavaSQLContext(self._jsc)
+
+    def sql(self, sqlQuery):
+        return SchemaRDD(self._jsql_ctx.sql(sqlQuery), self)
+
+    def applySchema(self, rdd):
+        jrdd = self._sc._pythonToJava(rdd._jrdd)
+        srdd = self._jsql_ctx.applySchema(jrdd)
+        return SchemaRDD(srdd, self)
+
+
 def _test():
     import atexit
     import doctest

python/pyspark/rdd.py

@@ -1389,13 +1389,17 @@ def _is_pipelinable(self):
 
 class SchemaRDD:
 
-    def __init__(self, pyRDD):
-        self._pyRDD = pyRDD
-        self.ctx = pyRDD.ctx
-        self.sql_ctx = self.ctx._jvm.JavaSQLContext(self.ctx._jsc)
-        self._jrdd = self.ctx._pythonToJava(pyRDD._jrdd)
-        self._srdd = self.sql_ctx.applySchema(self._jrdd)
+    def __init__(self, jschema_rdd, sql_ctx):
+        self.sql_ctx = sql_ctx
+        self._sc = sql_ctx._sc
+        self._jschema_rdd = jschema_rdd
 
+    def registerAsTable(self, name):
+        self._jschema_rdd.registerAsTable(name)
+
+    def toPython(self):
+        jrdd = self._sc._javaToPython(self._jschema_rdd)
+        return RDD(jrdd, self._sc, self._sc.serializer)
 
 def _test():
     import doctest