Merge pull request amplab#211 from alig/master

Improve the RDD <-> Table conversion support

Author: rxin

src/main/resources/tablerdd/RDDTable_generator.py

@@ -0,0 +1,95 @@
+#!/usr/bin/python
+from string import Template
+import sys
+from generator_utils import *
+
+## This script generates RDDtable.scala
+
+p = sys.stdout
+
+# e.g. createList(1,3, "T[", "]", ",") gives T[1],T[2],T[3]
+def createList(start, stop, prefix, suffix="", sep = ",", newlineAfter = 70, indent = 0):
+    res = ""
+    oneLine = res
+    for y in range(start,stop+1):
+        res     += prefix + str(y) + suffix
+        oneLine += prefix + str(y) + suffix
+        if y != stop:
+            res     += sep
+            oneLine += sep
+            if len(oneLine) > newlineAfter:
+                res += "\n" + " "*indent
+                oneLine = ""
+    return res
+
+### The SparkContext declaration
+
+prefix = """
+/*
+ * Copyright (C) 2012 The Regents of The University California.
+ * All rights reserved.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package shark.api
+
+// *** This file is auto-generated from RDDTable_generator.py ***
+
+import org.apache.spark.rdd.RDD
+
+object RDDTableImplicits {
+  private type M[T] = ClassManifest[T]
+
+"""
+
+p.write(prefix)
+
+for x in range(2,23):
+
+    tableClass = Template(
+"""
+  implicit def rddToTable$num[$tmlist]
+  (rdd: RDD[($tlist)]): RDDTableFunctions = RDDTable(rdd)
+
+""").substitute(num = x, tmlist = createList(1, x, "T", ": M", ", ", indent=4), tlist = createList(1, x, "T", "", ", ", indent=4))
+    p.write(tableClass)
+
+prefix = """
+}
+
+object RDDTable {
+
+  private type M[T] = ClassManifest[T]
+  private def m[T](implicit m : ClassManifest[T]) = classManifest[T](m)
+"""
+
+p.write(prefix)
+  
+for x in range(2,23):
+
+    tableClass = Template(
+"""
+  def apply[$tmlist]
+  (rdd: RDD[($tlist)]) = {
+    val cm = implicitly[Manifest[Seq[Any]]]
+    val rddSeq: RDD[Seq[_]] = rdd.map(t => t.productIterator.toList.asInstanceOf[Seq[Any]])(cm)
+    new RDDTableFunctions(rddSeq, Seq($mtlist))
+  }
+
+""").substitute(tmlist = createList(1, x, "T", ": M", ", ", indent=4), tlist = createList(1, x, "T", "", ", ", indent=4),
+                mtlist = createList(1, x, "m[T", "]", ", ", indent=4))
+    p.write(tableClass)
+
+
+p.write("}\n")

src/main/resources/tablerdd/SharkContext_sqlRdd_generator.py

@@ -0,0 +1,24 @@
+#!/usr/bin/python
+from string import Template
+import sys
+
+from generator_utils import *
+
+## This script generates functions sqlRdd for SharkContext.scala
+
+p = sys.stdout
+
+# The SharkContext declarations
+for x in range(2,23):
+    sqlRddFun = Template(
+"""
+  def sqlRdd[$list1](cmd: String):
+    RDD[Tuple$num[$list2]] = {
+    new TableRDD$num[$list2](sql2rdd(cmd),
+      Seq($list3))
+  }
+""").substitute(num = x, 
+               list1 = createList(1, x, "T", ": M", ", ", 80, 4), 
+               list2 = createList(1, x, "T", sep=", ", indent = 4), 
+               list3 = createList(1, x, "m[T", "]", sep=", ", indent = 10))
+    p.write(sqlRddFun)

src/main/resources/tablerdd/TableRDDGenerated_generator.py

@@ -0,0 +1,89 @@
+#!/usr/bin/python
+from string import Template
+import sys
+from generator_utils import *
+
+## This script generates TableRDDGenerated.scala
+
+p = sys.stdout
+
+p.write(
+"""
+/*
+ * Copyright (C) 2013 The Regents of The University California.
+ * All rights reserved.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+
+
+package shark.api
+
+// *** This file is auto-generated from TableRDDGenerated_generator.py ***
+
+import org.apache.spark.rdd.RDD
+import org.apache.spark.{TaskContext, Partition}
+
+class TableSeqRDD(prev: TableRDD)
+  extends RDD[Seq[Any]](prev) {
+
+  def getSchema = prev.schema
+
+  override def getPartitions = prev.getPartitions
+
+  override def compute(split: Partition, context: TaskContext): Iterator[Seq[Any]] = {
+    prev.compute(split, context).map( row =>
+      (0 until prev.schema.size).map(i => row.getPrimitive(i)) )
+  }
+}
+
+""")
+
+for x in range(1,23):
+
+    inner = ""
+    for y in range(1,x+1):
+        if y % 3 == 1: inner += "       "
+        inner += Template(" row.getPrimitiveGeneric[T$num1]($num2)").substitute(num1=y, num2=y-1)
+        if y != x: inner += ","
+        if y % 3 == 0: inner += "\n"
+    inner += " ) )\n"
+
+    tableClass = Template(
+"""
+class TableRDD$num[$list](prev: TableRDD,
+                       mans: Seq[ClassManifest[_]])
+  extends RDD[Tuple$num[$list]](prev) {
+  def schema = prev.schema
+
+  private val tableCols = schema.size
+  require(tableCols == $num, "Table only has " + tableCols + " columns, expecting $num")
+
+  mans.zipWithIndex.foreach{ case (m, i) => if (DataTypes.fromManifest(m) != schema(i).dataType)
+    throw new IllegalArgumentException(
+      "Type mismatch on column " + (i + 1) + ", expected " + DataTypes.fromManifest(m) + " got " + schema(i).dataType) }
+
+  override def getPartitions = prev.getPartitions
+
+  override def compute(split: Partition, context: TaskContext): 
+  Iterator[Tuple$num[$list]] = {
+    prev.compute(split, context).map( row =>
+      new Tuple$num[$list](
+        $innerfatlist
+  }
+}
+""").substitute(num = x, list = createList(1, x, "T", "", ", ", indent=4), innerfatlist = inner)
+
+
+    p.write(tableClass)

src/main/resources/tablerdd/generator_utils.py

@@ -0,0 +1,18 @@
+#!/usr/bin/python
+import sys
+
+# e.g. createList(1,3, "T[", "]", ",") gives T[1],T[2],T[3]
+def createList(start, stop, prefix, suffix="", sep = ",", newlineAfter = 70, indent = 0):
+    res = ""
+    oneLine = res
+    for y in range(start,stop+1):
+        res     += prefix + str(y) + suffix
+        oneLine += prefix + str(y) + suffix
+        if y != stop:
+            res     += sep
+            oneLine += sep
+            if len(oneLine) > newlineAfter:
+                res += "\n" + " "*indent
+                oneLine = ""
+    return res
+