Support IF in JS

hbutani · Sep 14, 2016 · ca31bac · ca31bac
1 parent 4f12e1d
commit ca31bac
Show file tree

Hide file tree

Showing 3 changed files with 43 additions and 6 deletions.
diff --git a/src/main/scala/org/sparklinedata/druid/jscodegen/JSCodeGenerator.scala b/src/main/scala/org/sparklinedata/druid/jscodegen/JSCodeGenerator.scala
@@ -96,7 +96,7 @@ case class JSCodeGenerator(dqb: DruidQueryBuilder, e: Expression, mulInParamsAll
                 dD.isInstanceOf[DruidTimeDimension])
               val cs = JSCast(jE, e.dataType, this).castCode
               cs.map( cs =>
-                JSExpr(cs.fnVar, jE.linesSoFar + cs.linesSoFar, cs.curLine, e.dataType)
+                JSExpr(cs.fnVar, jE.linesSoFar + cs.linesSoFar, cs.getRef, e.dataType)
               )
             }
           }.get
@@ -191,6 +191,19 @@ case class JSCodeGenerator(dqb: DruidQueryBuilder, e: Expression, mulInParamsAll
             None
           }
         }
+        case If(pe, te, fe) =>
+          for (
+            jspe <- genExprCode(pe);
+            jste <- genExprCode(te);
+            jsfe <- genExprCode(fe)) yield {
+            val v1 = makeUniqueVarName
+            JSExpr(Some(v1),
+              jspe.linesSoFar + jste.linesSoFar + jsfe.linesSoFar +
+              s"""var ${v1} = null;
+                 |if (${jspe.getRef})
+                 |{$v1 = ${jste.getRef}} else {$v1 = ${jsfe.getRef}}""".stripMargin,
+              "", jste.fnDT, false)
+          }
         case LessThan(l, r) => genComparisonCode(l, r, " < ")
         case LessThanOrEqual(l, r) => genComparisonCode(l, r, " <= ")
         case GreaterThan(l, r) => genComparisonCode(l, r, " > ")
@@ -308,7 +321,7 @@ case class JSCodeGenerator(dqb: DruidQueryBuilder, e: Expression, mulInParamsAll
         case DateFormatClass(de, fmt) => {
           for (jsDe <- genExprCode(de);
                jsFmt <- genExprCode(fmt);
-               fmtDStr <- dateFormatCode(jsDe, jsFmt.curLine)) yield {
+               fmtDStr <- dateFormatCode(jsDe, jsFmt.getRef)) yield {
             JSExpr(None, jsDe.linesSoFar + jsFmt.linesSoFar, fmtDStr, StringType, false)
           }
         }
@@ -456,7 +469,7 @@ case class JSCodeGenerator(dqb: DruidQueryBuilder, e: Expression, mulInParamsAll
   private[this] def genCastExprCode(e: Expression, dt: DataType): Option[JSExpr] = {
     for (fn <- genExprCode(ExprUtil.simplifyCast(e, dt));
          cs <- JSCast(fn, dt, this).castCode) yield
-      JSExpr(cs.fnVar, fn.linesSoFar + cs.linesSoFar, cs.curLine, dt)
+      JSExpr(cs.fnVar, fn.linesSoFar + cs.linesSoFar, cs.getRef, dt)
   }
 
   private[this] def validInParams(inParam: String): Boolean = {

diff --git a/src/main/scala/org/sparklinedata/druid/jscodegen/JSDateTime.scala b/src/main/scala/org/sparklinedata/druid/jscodegen/JSDateTime.scala
@@ -55,9 +55,9 @@ private[jscodegen] object JSDateTimeCtx {
 
   private[jscodegen] def dateFormatCode(dtVal: JSExpr, fmt: String): Option[String] =
     for (javaDate <- dtVal.fnDT match {
-      case TimestampType => Some(s"""${dtVal.curLine}.toDate()""".stripMargin)
-      case DateType => Some(s"""${dtVal.curLine}.toDate()""".stripMargin)
-      case StringType => Some(s"""${dtVal.curLine}""".stripMargin)
+      case TimestampType => Some(s"""${dtVal.getRef}.toDate()""".stripMargin)
+      case DateType => Some(s"""${dtVal.getRef}.toDate()""".stripMargin)
+      case StringType => Some(s"""${dtVal.getRef}""".stripMargin)
       case _ => None
     }) yield {
       s"(new java.text.SimpleDateFormat(${fmt})).format(${javaDate})"

diff --git a/src/test/scala/org/sparklinedata/druid/client/CodeGenTest.scala b/src/test/scala/org/sparklinedata/druid/client/CodeGenTest.scala
@@ -817,5 +817,29 @@ class CodeGenTest extends BaseTest with BeforeAndAfterAll with Logging {
       |order by x, y
     """.stripMargin
     ,1,true,true)
+
+  test("if1",
+    """
+      |select o_orderstatus as a,
+      |date_format(cast(o_orderdate as date), 'u') as b,
+      |if(date_format(cast(o_orderdate as date), 'u') = 1, 'true', 'false') as c
+      |from orderLineItemPartSupplier
+      |group by o_orderstatus, date_format(cast(o_orderdate as date), 'u'),
+      |if(date_format(cast(o_orderdate as date), 'u') = 1, 'true', 'false')
+      |order by a, b, c
+    """.stripMargin
+    ,1,true,true)
+
+  test("if2",
+    """
+      |select o_orderstatus as a,
+      |date_format(cast(o_orderdate as date), 'yy') as b,
+      |if(date_format(cast(o_orderdate as date), 'yy') = 94, 'true', 'false') as c
+      |from orderLineItemPartSupplier
+      |group by o_orderstatus, date_format(cast(o_orderdate as date), 'yy'),
+      |if(date_format(cast(o_orderdate as date), 'yy') = 94, 'true', 'false')
+      |order by a, b, c
+    """.stripMargin
+    ,1,true,true)
 }