DEV: Create class for returning union of features and covariates (#237)

aehrc · Sep 19, 2024 · 9733844 · 9733844
1 parent dfae3c2
commit 9733844
Showing 1 changed file with 15 additions and 0 deletions.
diff --git a/src/main/scala/au/csiro/variantspark/input/UnionedFeatureSource.scala b/src/main/scala/au/csiro/variantspark/input/UnionedFeatureSource.scala
@@ -0,0 +1,15 @@
+package au.csiro.variantspark.input
+
+import au.csiro.variantspark.data.Feature
+import org.apache.spark.rdd.RDD
+
+case class UnionedFeatureSource(featureSource: FeatureSource, covariateSource: FeatureSource)
+    extends FeatureSource {
+  def sampleNames: List[String] = featureSource.sampleNames
+
+  def features: RDD[Feature] = {
+    val features = featureSource.features
+    val covariates = covariateSource.features
+    features.union(covariates)
+  }
+}