Merge distinct agg (#192)

* Support for multiple branched CaseWhen

* Interval (#116)

* add date_add, interval sql still running into issues

* Add Interval SQL support

* uncomment out the other tests

* resolve comments

* change interval equality

Co-authored-by: Eric Feng <[email protected]>

* Remove partition ID argument from enclaves

* Fix comments

* updates

* Modifications to integrate crumb, log-mac, and all-outputs_mac, wip

* Store log mac after each output buffer, add all-outputs-mac to each encryptedblocks wip

* Add all_outputs_mac to all EncryptedBlocks once all log_macs have been generated

* Almost builds

* cpp builds

* Use ubyte for all_outputs_mac

* use Mac for all_outputs_mac

* Hopefully this works for flatbuffers all_outputs_mac mutation, cpp builds

* Scala builds now too, running into error with union

* Stuff builds, error with all outputs mac serialization. this commit uses all_outputs_mac as Mac table

* Fixed bug, basic encryption / show works

* All single partition tests pass, multiple partiton passes until tpch-9

* All tests pass except tpch-9 and skew join

* comment tpch back in

* Check same number of ecalls per partition - exception for scanCollectLastPrimary(?)

* First attempt at constructing executed DAG

* Fix typos

* Rework graph

* Add log macs to graph nodes

* Construct expected DAG and refactor JobNode.
Refactor construction of executed DAG.

* Implement 'paths to sink' for a DAG

* add crumb for last ecall

* Fix NULL handling for aggregation (#130)

* Modify COUNT and SUM to correctly handle NULL values

* Change average to support NULL values

* Fix

* Changing operator matching from logical to physical (#129)

* WIP

* Fix

* Unapply change

* Aggregation rewrite (#132)

* updated build/sbt file (#135)

* Travis update (#137)

* update breeze (#138)

* TPC-H test suite added (#136)

* added tpch sql files

* functions updated to save temp view

* main function skeleton done

* load and clear done

* fix clear

* performQuery done

* import cleanup, use OPAQUE_HOME

* TPC-H 9 refactored to use SQL rather than DF operations

* removed : Unit, unused imports

* added TestUtils.scala

* moved all common initialization to TestUtils

* update name

* begin rewriting TPCH.scala to store persistent tables

* invalid table name error

* TPCH conversion to class started

* compiles

* added second case, cleared up names

* added TPC-H 6 to check that persistent state has no issues

* added functions for the last two tables

* addressed most logic changes

* DataFrame only loaded once

* apply method in companion object

* full test suite added

* added testFunc parameter to testAgainstSpark

* ignore #18

* Separate IN PR (#124)

* finishing the in expression. adding more tests and null support. need confirmation on null behavior and also I wonder why integer field is sufficient for string

* adding additional test

* adding additional test

* saving concat implementation and it's passing basic functionality tests

* adding type aware comparison and better error message for IN operator

* adding null checking for the concat operator and adding one additional test

* cleaning up IN&Concat PR

* deleting concat and preping the in branch for in pr

* fixing null bahavior 

now it's only null when there's no match and there's null input

* Build failed

Co-authored-by: Ubuntu <[email protected]>
Co-authored-by: Wenting Zheng <[email protected]>
Co-authored-by: Wenting Zheng <[email protected]>

* Merge new aggregate

* Uncomment log_mac_lst clear

* Clean up comments

* Separate Concat PR  (#125)

Implementation of the CONCAT expression.

Co-authored-by: Ubuntu <[email protected]>
Co-authored-by: Wenting Zheng <[email protected]>

* Clean up comments in other files

* Update pathsEqual to be less conservative

* Remove print statements from unit tests

* Removed calls to toSet in TPC-H tests (#140)

* removed calls to toSet

* added calls to toSet back where queries are unordered

* Documentation update (#148)

* Cluster Remote Attestation Fix (#146)

The existing code only had RA working when run locally. This PR adds a sleep for 5 seconds to make sure that all executors are spun up successfully before attestation begins.

Closes #147

* upgrade to 3.0.1 (#144)

* Update two TPC-H queries (#149)

Tests for TPC-H 12 and 19 pass.

* TPC-H 20 Fix (#142)

* string to stringtype error

* tpch 20 passes

* cleanup

* implemented changes

* decimal.tofloat

Co-authored-by: Wenting Zheng <[email protected]>

* Add expected operator DAG generation from executedPlan string

* Rebase

* Join update (#145)

* Merge join update

* Integrate new join

* Add expected operator for sortexec

* Merge comp-integrity with join update

* Remove some print statements

* Migrate from Travis CI to Github Actions (#156)

* Upgrade to OE 0.12 (#153)

* Update README.md

* Support for scalar subquery (#157)

This PR implements the scalar subquery expression, which is triggered whenever a subquery returns a scalar value. There were two main problems that needed to be solved.

First, support for matching the scalar subquery expression is necessary. Spark implements this by wrapping a SparkPlan within the expression and calls executeCollect. Then it constructs a literal with that value. However, this is problematic for us because that value should not be decrypted by the driver and serialized into an expression, since it's an intermediate value.

Therefore, the second issue to be addressed here is supporting an encrypted literal. This is implemented in this PR by serializing an encrypted ciphertext into a base64 encoded string, and wrapping a Decrypt expression on top of it. This expression is then evaluated in the enclave and returns a literal. Note that, in order to test our implementation, we also implement a Decrypt expression in Scala. However, this should never be evaluated on the driver side and serialized into a plaintext literal. This is because Decrypt is designated as a Nondeterministic expression, and therefore will always evaluate on the workers.

* Add TPC-H Benchmarks (#139)

* logic decoupling in TPCH.scala for easier benchmarking

* added TPCHBenchmark.scala

* Benchmark.scala rewrite

* done adding all support TPC-H query benchmarks

* changed commandline arguments that benchmark takes

* TPCHBenchmark takes in parameters

* fixed issue with spark conf

* size error handling, --help flag

* add Utils.force, break cluster mode

* comment out logistic regression benchmark

* ensureCached right before temp view created/replaced

* upgrade to 3.0.1

* upgrade to 3.0.1

* 10 scale factor

* persistData

* almost done refactor

* more cleanup

* compiles

* 9 passes

* cleanup

* collect instead of force, sf_none

* remove sf_none

* defaultParallelism

* no removing trailing/leading whitespace

* add sf_med

* hdfs works in local case

* cleanup, added new CLI argument

* added newly supported tpch queries

* function for running all supported tests

* Construct expected DAG from dataframe physical plan

* Refactor collect and add integrity checking helper function to OpaqueOperatorTest

* Float expressions (#160)

This PR adds float normalization expressions [implemented in Spark](https://github.com/apache/spark/blob/master/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/NormalizeFloatingNumbers.scala#L170). TPC-H query 2 also passes.

* Broadcast Nested Loop Join - Left Anti and Left Semi  (#159)

This PR is the first of two parts towards making TPC-H 16 work: the other will be implementing `is_distinct` for aggregate operations.

`BroadcastNestedLoopJoin` is Spark's "catch all" for non-equi joins. It works by first picking a side to broadcast, then iterating through every possible row combination and checking the non-equi condition against the pair.

* Move join condition handling for equi-joins into enclave code (#164)

* Add in TPC-H 21

* Add condition processing in enclave code

* Code clean up

* Enable query 18

* WIP

* Local tests pass

* Apply suggestions from code review

Co-authored-by: octaviansima <[email protected]>

* WIP

* Address comments

* q21.sql

Co-authored-by: octaviansima <[email protected]>

* Distinct aggregation support (#163)

* matching in strategies.scala

set up class thing

cleanup

added test cases for non-equi left anti join

rename to serializeEquiJoinExpression

added isEncrypted condition

set up keys

JoinExpr now has condition

rename

serialization does not throw compile error for BNLJ

split up

added condition in ExpressionEvaluation.h

zipPartitions

cpp put in place

typo

added func to header

two loops in place

update tests

condition

fixed scala loop

interchange rows

added tags

ensure cached

== match working

comparison decoupling in ExpressionEvalulation

save

compiles and condition works

is printing

fix swap outer/inner

o_i_match

show() has the same result

tests pass

test cleanup

added test cases for different condition

BuildLeft works

optional keys in scala

started C++

passes the operator tests

comments, cleanup

attemping to do it the ~right~ way

comments to distinguish between primary/secondary, operator tests pass

cleanup comments, about to begin implementation for distinct agg ops

is_distinct

added test case

serializing with isDistinct

is_distinct in ExpressionEvaluation.h

removed unused code from join implementation

remove RowWriter/Reader in condition evaluation (join)

easier test

serialization done

correct checking in Scala

set is set up

spaghetti but it finally works

function for clearing values

condition_eval isntead of condition

goto

comment

remove explain from test, need to fix distinct aggregation for >1 partitions

started impl of multiple partitions fix

added rangepartitionexec that runs

partitioning cleanup

serialization properly

comments, generalization for > 1 distinct function

comments

about to refactor into logical.Aggregation

the new case has distinct in result expressions

need to match on distinct

removed new case (doesn't make difference?)

works

Upgrade to OE 0.12 (#153)

Update README.md

Support for scalar subquery (#157)

This PR implements the scalar subquery expression, which is triggered whenever a subquery returns a scalar value. There were two main problems that needed to be solved.

First, support for matching the scalar subquery expression is necessary. Spark implements this by wrapping a SparkPlan within the expression and calls executeCollect. Then it constructs a literal with that value. However, this is problematic for us because that value should not be decrypted by the driver and serialized into an expression, since it's an intermediate value.

Therefore, the second issue to be addressed here is supporting an encrypted literal. This is implemented in this PR by serializing an encrypted ciphertext into a base64 encoded string, and wrapping a Decrypt expression on top of it. This expression is then evaluated in the enclave and returns a literal. Note that, in order to test our implementation, we also implement a Decrypt expression in Scala. However, this should never be evaluated on the driver side and serialized into a plaintext literal. This is because Decrypt is designated as a Nondeterministic expression, and therefore will always evaluate on the workers.

match

remove RangePartitionExec

inefficient implementation refined

Add TPC-H Benchmarks (#139)

* logic decoupling in TPCH.scala for easier benchmarking

* added TPCHBenchmark.scala

* Benchmark.scala rewrite

* done adding all support TPC-H query benchmarks

* changed commandline arguments that benchmark takes

* TPCHBenchmark takes in parameters

* fixed issue with spark conf

* size error handling, --help flag

* add Utils.force, break cluster mode

* comment out logistic regression benchmark

* ensureCached right before temp view created/replaced

* upgrade to 3.0.1

* upgrade to 3.0.1

* 10 scale factor

* persistData

* almost done refactor

* more cleanup

* compiles

* 9 passes

* cleanup

* collect instead of force, sf_none

* remove sf_none

* defaultParallelism

* no removing trailing/leading whitespace

* add sf_med

* hdfs works in local case

* cleanup, added new CLI argument

* added newly supported tpch queries

* function for running all supported tests

complete instead of partial -> final

removed traces of join

cleanup

* added test case for one distinct one non, reverted comment

* removed C++ level implementation of is_distinct

* PartialMerge in operators.scala

* stage 1: grouping with distinct expressions

* stage 2: WIP

* saving, sorting by group expressions ++ name distinct expressions worked

* stage 1 & 2 printing the expected results

* removed extraneous call to sorted, #3 in place but not working

* stage 3 has the final, correct result: refactoring the Aggregate code to not cast aggregate expressions to Partial, PartialMerge, etc will be needed

* refactor done, C++ still printing the correct values

* need to formalize None case in EncryptedAggregateExec.output, but stage 4 passes

* distinct and indistinct passes (git add -u)

* general cleanup, None case looks nicer

* throw error with >1 distinct, add test case for global distinct

* no need for global aggregation case

* single partition passes all aggregate tests, multiple partition doesn't

* works with global sort first

* works with non-global sort first

* cleanup

* cleanup tests

* removed iostream, other nit

* added test case for 13

* None case in isPartial match done properly

* added test cases for sumDistinct

* case-specific namedDistinctExpressions working

* distinct sum is done

* removed comments

* got rid of mode argument

* tests include null values

* partition followed by local sort instead of first global sort

* Remove addExpectedOperator from JobVerificationEngine, add comments

* Implement expected DAG construction by doing graph manipulation on dataframe field instead of string parsing

* Fix merge errors in the test cases

Co-authored-by: Andrew Law <[email protected]>
Co-authored-by: Eric Feng <[email protected]>
Co-authored-by: Eric Feng <[email protected]>
Co-authored-by: Chester Leung <[email protected]>
Co-authored-by: Wenting Zheng <[email protected]>
Co-authored-by: octaviansima <[email protected]>
Co-authored-by: Chenyu Shi <[email protected]>
Co-authored-by: Ubuntu <[email protected]>
Co-authored-by: Wenting Zheng <[email protected]>

Author: 9 people

src/main/scala/edu/berkeley/cs/rise/opaque/Utils.scala

@@ -1435,6 +1435,10 @@ object Utils extends Logging {
              }
             (Seq(countUpdateExpr), Seq(count))
           }
+          case PartialMerge => {
+            val countUpdateExpr = Add(count, c.inputAggBufferAttributes(0))
+            (Seq(countUpdateExpr), Seq(count))
+          }
           case Final => {
             val countUpdateExpr = Add(count, c.inputAggBufferAttributes(0))
             (Seq(countUpdateExpr), Seq(count))
@@ -1443,7 +1447,7 @@ object Utils extends Logging {
             val countUpdateExpr = Add(count, Literal(1L))
             (Seq(countUpdateExpr), Seq(count))
           }
-          case _ => 
+          case _ =>
         }
 
         tuix.AggregateExpr.createAggregateExpr(
@@ -1614,6 +1618,11 @@ object Utils extends Logging {
             val sumUpdateExpr = If(IsNull(partialSum), sum, partialSum)
             (Seq(sumUpdateExpr), Seq(sum))
           }
+          case PartialMerge => {
+            val partialSum = Add(If(IsNull(sum), Literal.default(sumDataType), sum), s.inputAggBufferAttributes(0))
+            val sumUpdateExpr = If(IsNull(partialSum), sum, partialSum)
+            (Seq(sumUpdateExpr), Seq(sum))
+          }
           case Final => {
             val partialSum = Add(If(IsNull(sum), Literal.default(sumDataType), sum), s.inputAggBufferAttributes(0))
             val sumUpdateExpr = If(IsNull(partialSum), sum, partialSum)

src/main/scala/edu/berkeley/cs/rise/opaque/benchmark/TPCHBenchmark.scala

@@ -24,7 +24,7 @@ import org.apache.spark.sql.SQLContext
 object TPCHBenchmark {
 
   // Add query numbers here once they are supported
-  val supportedQueries = Seq(1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 14, 15, 17, 18, 19, 20, 21, 22)
+  val supportedQueries = Seq(1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 14, 15, 16, 17, 18, 19, 20, 21, 22)
 
   def query(queryNumber: Int, tpch: TPCH, sqlContext: SQLContext, numPartitions: Int) = {
     val sqlStr = tpch.getQuery(queryNumber)

src/main/scala/edu/berkeley/cs/rise/opaque/execution/operators.scala

@@ -243,44 +243,35 @@ case class EncryptedFilterExec(condition: Expression, child: SparkPlan)
 
 case class EncryptedAggregateExec(
   groupingExpressions: Seq[NamedExpression],
-  aggExpressions: Seq[AggregateExpression],
-  mode: AggregateMode,
+  aggregateExpressions: Seq[AggregateExpression],
   child: SparkPlan)
     extends UnaryExecNode with OpaqueOperatorExec {
 
   override def producedAttributes: AttributeSet =
-    AttributeSet(aggExpressions) -- AttributeSet(groupingExpressions)
-
-  override def output: Seq[Attribute] = mode match {
-    case Partial => groupingExpressions.map(_.toAttribute) ++ aggExpressions.map(_.copy(mode = Partial)).flatMap(_.aggregateFunction.inputAggBufferAttributes)
-    case Final => groupingExpressions.map(_.toAttribute) ++ aggExpressions.map(_.resultAttribute)
-    case Complete => groupingExpressions.map(_.toAttribute) ++ aggExpressions.map(_.resultAttribute)
-  }
+    AttributeSet(aggregateExpressions) -- AttributeSet(groupingExpressions)
+
+  override def output: Seq[Attribute] = groupingExpressions.map(_.toAttribute) ++
+    aggregateExpressions.flatMap(expr => {
+      expr.mode match {
+        case Partial | PartialMerge =>
+          expr.aggregateFunction.inputAggBufferAttributes
+        case _ =>
+          Seq(expr.resultAttribute)
+      }
+    })
 
   override def executeBlocked(): RDD[Block] = {
 
-    val (groupingExprs, aggExprs) = mode match {
-      case Partial => {
-        val partialAggExpressions = aggExpressions.map(_.copy(mode = Partial))
-        (groupingExpressions, partialAggExpressions)
-      }
-      case Final => {
-        val finalGroupingExpressions = groupingExpressions.map(_.toAttribute)
-        val finalAggExpressions = aggExpressions.map(_.copy(mode = Final))
-        (finalGroupingExpressions, finalAggExpressions)
-      }
-      case Complete => {
-        (groupingExpressions, aggExpressions.map(_.copy(mode = Complete)))
-      }
-    }
+    val aggExprSer = Utils.serializeAggOp(groupingExpressions, aggregateExpressions, child.output)
+    val isPartial = aggregateExpressions.map(expr => expr.mode)
+      .exists(mode => mode == Partial || mode == PartialMerge)
 
-     val aggExprSer = Utils.serializeAggOp(groupingExprs, aggExprs, child.output)
 
     timeOperator(child.asInstanceOf[OpaqueOperatorExec].executeBlocked(), "EncryptedPartialAggregateExec") {
       childRDD => 
         childRDD.map { block =>
         val (enclave, eid) = Utils.initEnclave()
-        Block(enclave.NonObliviousAggregate(eid, aggExprSer, block.bytes, (mode == Partial)))
+        Block(enclave.NonObliviousAggregate(eid, aggExprSer, block.bytes, isPartial))
       }
     }
   }

src/main/scala/edu/berkeley/cs/rise/opaque/strategies.scala

@@ -152,25 +152,90 @@ object OpaqueOperators extends Strategy {
         if (isEncrypted(child) && aggExpressions.forall(expr => expr.isInstanceOf[AggregateExpression])) =>
 
       val aggregateExpressions = aggExpressions.map(expr => expr.asInstanceOf[AggregateExpression])
-
-      if (groupingExpressions.size == 0) {
-        // Global aggregation
-        val partialAggregate = EncryptedAggregateExec(groupingExpressions, aggregateExpressions, Partial, planLater(child))
-        val partialOutput = partialAggregate.output
-        val (projSchema, tag) = tagForGlobalAggregate(partialOutput)
-
-        EncryptedProjectExec(resultExpressions, 
-          EncryptedAggregateExec(groupingExpressions, aggregateExpressions, Final, 
-            EncryptedProjectExec(partialOutput, 
-              EncryptedSortExec(Seq(SortOrder(tag, Ascending)), true, 
-                EncryptedProjectExec(projSchema, partialAggregate))))) :: Nil
-      } else {
-        // Grouping aggregation
-        EncryptedProjectExec(resultExpressions,
-          EncryptedAggregateExec(groupingExpressions, aggregateExpressions, Final,
-            EncryptedSortExec(groupingExpressions.map(_.toAttribute).map(e => SortOrder(e, Ascending)), true,
-              EncryptedAggregateExec(groupingExpressions, aggregateExpressions, Partial,
-                EncryptedSortExec(groupingExpressions.map(e => SortOrder(e, Ascending)), false, planLater(child)))))) :: Nil
+      val (functionsWithDistinct, functionsWithoutDistinct) = aggregateExpressions.partition(_.isDistinct)
+
+      functionsWithDistinct.size match {
+        case 0 => // No distinct aggregate operations
+          if (groupingExpressions.size == 0) {
+            // Global aggregation
+            val partialAggregate = EncryptedAggregateExec(groupingExpressions,
+              aggregateExpressions.map(_.copy(mode = Partial)), planLater(child))
+            val partialOutput = partialAggregate.output
+            val (projSchema, tag) = tagForGlobalAggregate(partialOutput)
+
+            EncryptedProjectExec(resultExpressions, 
+              EncryptedAggregateExec(groupingExpressions, aggregateExpressions.map(_.copy(mode = Final)),
+                EncryptedProjectExec(partialOutput,
+                  EncryptedSortExec(Seq(SortOrder(tag, Ascending)), true, 
+                    EncryptedProjectExec(projSchema, partialAggregate))))) :: Nil
+          } else {
+            // Grouping aggregation
+            EncryptedProjectExec(resultExpressions,
+              EncryptedAggregateExec(groupingExpressions, aggregateExpressions.map(_.copy(mode = Final)),
+                EncryptedSortExec(groupingExpressions.map(_.toAttribute).map(e => SortOrder(e, Ascending)), true,
+                  EncryptedAggregateExec(groupingExpressions, aggregateExpressions.map(_.copy(mode = Partial)),
+                    EncryptedSortExec(groupingExpressions.map(e => SortOrder(e, Ascending)), false, planLater(child)))))) :: Nil
+          }
+        case size if size == 1 => // One distinct aggregate operation
+          // Because we are also grouping on the columns used in the distinct expressions,
+          // we do not need separate cases for global and grouping aggregation.
+
+          // We need to extract named expressions from the children of the distinct aggregate functions
+          // in order to group by those columns.
+          val namedDistinctExpressions = functionsWithDistinct.head.aggregateFunction.children.flatMap{ e =>
+            e match {
+              case ne: NamedExpression =>
+                Seq(ne)
+              case _ =>
+                e.children.filter(child => child.isInstanceOf[NamedExpression])
+                  .map(child => child.asInstanceOf[NamedExpression])
+            }
+          }
+          val combinedGroupingExpressions = groupingExpressions ++ namedDistinctExpressions
+
+          // 1. Create an Aggregate operator for partial aggregations.
+          val partialAggregate = {
+            val sorted = EncryptedSortExec(combinedGroupingExpressions.map(e => SortOrder(e, Ascending)), false,
+              planLater(child))
+            EncryptedAggregateExec(combinedGroupingExpressions, functionsWithoutDistinct.map(_.copy(mode = Partial)), sorted)
+          }
+
+          // 2. Create an Aggregate operator for partial merge aggregations.
+          val partialMergeAggregate = {
+            // Partition based on the final grouping expressions.
+            val partitionOrder = groupingExpressions.map(e => SortOrder(e, Ascending))
+            val partitioned = EncryptedRangePartitionExec(partitionOrder, partialAggregate)
+
+            // Local sort on the combined grouping expressions.
+            val sortOrder = combinedGroupingExpressions.map(e => SortOrder(e, Ascending))
+            val sorted = EncryptedSortExec(sortOrder, false, partitioned)
+
+            EncryptedAggregateExec(combinedGroupingExpressions,
+              functionsWithoutDistinct.map(_.copy(mode = PartialMerge)), sorted)
+          }
+
+          // 3. Create an Aggregate operator for partial aggregation of distinct aggregate expressions.
+          val partialDistinctAggregate = {
+            // Indistinct functions operate on aggregation buffers since partial aggregation was already called,
+            // but distinct functions operate on the original input to the aggregation.
+            EncryptedAggregateExec(groupingExpressions,
+              functionsWithoutDistinct.map(_.copy(mode = PartialMerge)) ++
+                functionsWithDistinct.map(_.copy(mode = Partial)), partialMergeAggregate)
+          }
+
+          // 4. Create an Aggregate operator for the final aggregation.
+          val finalAggregate = {
+            val sorted = EncryptedSortExec(groupingExpressions.map(e => SortOrder(e, Ascending)),
+              true, partialDistinctAggregate)
+            EncryptedAggregateExec(groupingExpressions,
+                (functionsWithoutDistinct ++ functionsWithDistinct).map(_.copy(mode = Final)), sorted)
+          }
+
+          EncryptedProjectExec(resultExpressions, finalAggregate) :: Nil
+
+        case _ => { // More than one distinct operations
+          throw new UnsupportedOperationException("Aggregate operations with more than one distinct expressions are not yet supported.")
+        }
       }
 
     case p @ Union(Seq(left, right)) if isEncrypted(p) =>

src/test/scala/edu/berkeley/cs/rise/opaque/OpaqueOperatorTests.scala

@@ -489,6 +489,30 @@ trait OpaqueOperatorTests extends OpaqueTestsBase { self =>
       .sortBy { case Row(category: String, _) => category }
   }
 
+  testAgainstSpark("aggregate count distinct and indistinct") { securityLevel =>
+    val data = (0 until 64).map{ i =>
+      if (i % 6 == 0)
+        (abc(i), null.asInstanceOf[Int], i % 8)
+      else
+        (abc(i), i % 4, i % 8)
+    }.toSeq
+    val words = makeDF(data, securityLevel, "category", "id", "price")
+    words.groupBy("category").agg(countDistinct("id").as("num_unique_ids"),
+        count("price").as("num_prices")).collect.toSet
+  }
+
+  testAgainstSpark("aggregate count distinct") { securityLevel =>
+    val data = (0 until 64).map{ i =>
+      if (i % 6 == 0)
+        (abc(i), null.asInstanceOf[Int])
+      else
+        (abc(i), i % 8)
+    }.toSeq
+    val words = makeDF(data, securityLevel, "category", "price")
+    words.groupBy("category").agg(countDistinct("price").as("num_unique_prices"))
+      .collect.sortBy { case Row(category: String, _) => category }
+  }
+
   testAgainstSpark("aggregate first") { securityLevel =>
     val data = for (i <- 0 until 256) yield (i, abc(i), 1)
     val words = makeDF(data, securityLevel, "id", "category", "price")
@@ -536,6 +560,30 @@ trait OpaqueOperatorTests extends OpaqueTestsBase { self =>
       .sortBy { case Row(word: String, _) => word }
   }
 
+  testAgainstSpark("aggregate sum distinct and indistinct") { securityLevel =>
+    val data = (0 until 64).map{ i =>
+      if (i % 6 == 0)
+        (abc(i), null.asInstanceOf[Int], i % 8)
+      else
+        (abc(i), i % 4, i % 8)
+    }.toSeq
+    val words = makeDF(data, securityLevel, "category", "id", "price")
+    words.groupBy("category").agg(sumDistinct("id").as("sum_unique_ids"),
+        sum("price").as("sum_prices")).collect.toSet
+  }
+
+  testAgainstSpark("aggregate sum distinct") { securityLevel =>
+    val data = (0 until 64).map{ i =>
+      if (i % 6 == 0)
+        (abc(i), null.asInstanceOf[Int])
+      else
+        (abc(i), i % 8)
+    }.toSeq
+    val words = makeDF(data, securityLevel, "category", "price")
+    words.groupBy("category").agg(sumDistinct("price").as("sum_unique_prices"))
+      .collect.sortBy { case Row(category: String, _) => category }
+  }
+
   testAgainstSpark("aggregate on multiple columns") { securityLevel =>
     val data = for (i <- 0 until 256) yield (abc(i), 1, 1.0f)
     val words = makeDF(data, securityLevel, "str", "x", "y")
@@ -567,6 +615,12 @@ trait OpaqueOperatorTests extends OpaqueTestsBase { self =>
     integrityCollect(words.agg(sum("count").as("totalCount")))
   }
 
+  testAgainstSpark("global aggregate count distinct") { securityLevel =>
+    val data = for (i <- 0 until 256) yield (i, abc(i), i % 64)
+    val words = makeDF(data, securityLevel, "id", "word", "price")
+    words.agg(countDistinct("price").as("num_unique_prices")).collect
+  }
+
   testAgainstSpark("global aggregate with 0 rows") { securityLevel =>
     val data = for (i <- 0 until 256) yield (i, abc(i), 1)
     val words = makeDF(data, securityLevel, "id", "word", "count")