Merge pull request #39 from VirtusLab/better-match-types

Better match types

Author: prolativ

src/main/ColumnOp.scala

@@ -4,7 +4,7 @@ import scala.quoted._
 import org.apache.spark.sql
 import org.apache.spark.sql.SparkSession
 import org.virtuslab.iskra.DataFrame
-import org.virtuslab.iskra.types.{DataType, StructType, Encoder, StructEncoder, PrimitiveEncoder}
+import org.virtuslab.iskra.types.{DataType, Encoder, StructEncoder, PrimitiveEncoder}
 
 object DataFrameBuilders:
   extension [A](seq: Seq[A])(using encoder: Encoder[A])

src/main/DataFrameBuilders.scala

@@ -6,20 +6,30 @@ import types.{DataType, Encoder, StructEncoder}
 import MacroHelpers.TupleSubtype
 
 object FrameSchema:
-  type Merge[S1, S2] = S1 match
-    case TupleSubtype[s1] => S2 match
-      case TupleSubtype[s2] => Tuple.Concat[s1, s2]
-      case _ => Tuple.Concat[s1, S2 *: EmptyTuple]
-    case _ => S2 match
-      case TupleSubtype[s2] => S1 *: s2
-      case _ => S1 *: S2 *: EmptyTuple
+  type AsTuple[A] = A match
+    case Tuple => A
+    case Any => A *: EmptyTuple
+
+  type FromTuple[T] = T match
+    case h *: EmptyTuple => h
+    case Tuple => T
+
+  type Merge[S1, S2] = (S1, S2) match
+    case (Tuple, Tuple) =>
+      Tuple.Concat[S1, S2]
+    case (Any, Tuple) =>
+      S1 *: S2
+    case (Tuple, Any) =>
+      Tuple.Append[S1, S2]
+    case (Any, Any) =>
+      (S1, S2)
 
   type NullableLabeledDataType[T] = T match
-    case label := tpe => label := DataType.Nullable[tpe]
+    case label := tpe => label := DataType.AsNullable[tpe]
 
   type NullableSchema[T] = T match
-    case TupleSubtype[s] => Tuple.Map[s, NullableLabeledDataType]
-    case _ => NullableLabeledDataType[T]
+    case Tuple => Tuple.Map[T, NullableLabeledDataType]
+    case Any => NullableLabeledDataType[T]
 
   def reownType[Owner <: Name : Type](schema: Type[?])(using Quotes): Type[?] =
     schema match

src/main/FrameSchema.scala

@@ -3,7 +3,7 @@ package org.virtuslab.iskra
 import scala.language.implicitConversions
 
 import scala.quoted.*
-import org.virtuslab.iskra.types.BooleanOptType
+import org.virtuslab.iskra.types.BooleanOptLike
 
 trait OnConditionJoiner[Join <: JoinType, Left, Right]
 
@@ -73,7 +73,7 @@ object JoinOnCondition:
     import quotes.reflect.*
 
     '{ ${ condition }(using ${ joiningView }) } match
-      case '{ $cond: Col[BooleanOptType] } =>
+      case '{ $cond: Col[BooleanOptLike] } =>
         '{
           val joined = ${ join }.left.join(${ join }.right, ${ cond }.untyped, JoinType.typeName[T])
           StructDataFrame[JoinedSchema](joined)

src/main/JoinOnCondition.scala

@@ -11,7 +11,3 @@ private[iskra] object MacroHelpers:
     Position(file, start, end)
 
   type TupleSubtype[T <: Tuple] = T
-
-  type AsTuple[A] <: Tuple = A match
-    case TupleSubtype[t] => t
-    case _ => A *: EmptyTuple

src/main/MacroHelpers.scala

@@ -3,7 +3,7 @@ package org.virtuslab.iskra
 import scala.quoted.*
 import org.apache.spark.sql.functions.col
 import types.DataType
-import MacroHelpers.AsTuple
+import MacroHelpers.TupleSubtype
 
 inline def $(using view: SchemaView): view.type = view
 
@@ -79,7 +79,7 @@ object StructSchemaView:
     import quotes.reflect.*
     Type.of[DF] match
       case '[StructDataFrame[schema]] =>
-        val schemaType = Type.of[AsTuple[schema]]
+        val schemaType = Type.of[FrameSchema.AsTuple[schema]]
         val aliasViewsByName = frameAliasViewsByName(schemaType)
         val columns = unambiguousColumns(schemaType)    
         val frameAliasNames = Expr(aliasViewsByName.map(_._1))

src/main/SchemaView.scala

@@ -37,13 +37,7 @@ object Select:
       case Some(collectColumns) =>
         collectColumns match
           case '{ $cc: CollectColumns[?] { type CollectedColumns = collectedColumns } } =>
-            Type.of[collectedColumns] match
-              case '[head *: EmptyTuple] =>
-                '{
-                  val cols = ${ cc }.underlyingColumns(${ columns }(using ${ select }.view))
-                  StructDataFrame[head](${ select }.underlying.select(cols*))
-                }
-
+            Type.of[FrameSchema.FromTuple[collectedColumns]] match
               case '[s] =>
                 '{
                   val cols = ${ cc }.underlyingColumns(${ columns }(using ${ select }.view))

src/main/Select.scala

@@ -3,6 +3,7 @@ package org.virtuslab.iskra
 import scala.quoted.*
 
 import types.{DataType, Encoder, StructEncoder}
+import MacroHelpers.TupleSubtype
 
 
 class StructDataFrame[Schema](val untyped: UntypedDataFrame) extends DataFrame
@@ -20,7 +21,13 @@ object StructDataFrame:
     Expr.summon[Encoder[A]] match
       case Some(encoder) => encoder match
         case '{ $enc: StructEncoder[A] { type StructSchema = structSchema } } =>
-          Type.of[MacroHelpers.AsTuple[FrameSchema]] match
+          val frameSchemaTuple = Type.of[FrameSchema] match
+            case '[TupleSubtype[t]] =>
+              Type.of[t]
+            case '[t] =>
+              Type.of[t *: EmptyTuple]
+
+          frameSchemaTuple match
             case '[`structSchema`] =>
               '{ ClassDataFrame[A](${ df }.untyped) }
             case _ =>

src/main/StructDataFrame.scala

@@ -1,7 +1,7 @@
 package org.virtuslab.iskra
 
 import scala.quoted.*
-import types.{DataType, Encoder, StructType, StructEncoder}
+import types.{DataType, Encoder, struct, StructEncoder, StructNotNull}
 
 object UntypedOps:
   extension (untyped: UntypedColumn)
@@ -12,5 +12,5 @@ object UntypedOps:
 
   private def typedDataFrameImpl[A : Type](df: Expr[UntypedDataFrame], encoder: Expr[StructEncoder[A]])(using Quotes) =
     encoder match
-      case '{ ${e}: Encoder.Aux[tpe, StructType[t]] } =>
+      case '{ ${e}: Encoder.Aux[tpe, StructNotNull[t]] } =>
         '{ ClassDataFrame[A](${ df }) }

src/main/UntypedOps.scala

@@ -1,14 +1,14 @@
 package org.virtuslab.iskra
 
 import org.apache.spark.sql.{functions => f, Column => UntypedColumn}
-import org.virtuslab.iskra.types.{Coerce, DataType, BooleanOptType}
+import org.virtuslab.iskra.types.{Coerce, DataType, BooleanOptLike}
 
 object When:
-  class WhenColumn[T <: DataType](untyped: UntypedColumn) extends Col[DataType.Nullable[T]](untyped):
-    def when[U <: DataType](condition: Col[BooleanOptType], value: Col[U])(using coerce: Coerce[T, U]): WhenColumn[coerce.Coerced] =
+  class WhenColumn[T <: DataType](untyped: UntypedColumn) extends Col[DataType.AsNullable[T]](untyped):
+    def when[U <: DataType](condition: Col[BooleanOptLike], value: Col[U])(using coerce: Coerce[T, U]): WhenColumn[coerce.Coerced] =
       WhenColumn(this.untyped.when(condition.untyped, value.untyped))
     def otherwise[U <: DataType](value: Col[U])(using coerce: Coerce[T, U]): Col[coerce.Coerced] =
       Col(this.untyped.otherwise(value.untyped))
 
-  def when[T <: DataType](condition: Col[BooleanOptType], value: Col[T]): WhenColumn[T] =
+  def when[T <: DataType](condition: Col[BooleanOptLike], value: Col[T]): WhenColumn[T] =
     WhenColumn(f.when(condition.untyped, value.untyped))

src/main/When.scala

@@ -1,7 +1,7 @@
 package org.virtuslab.iskra
 
 import scala.quoted.*
-import org.virtuslab.iskra.types.BooleanOptType
+import org.virtuslab.iskra.types.BooleanOptLike
 
 trait Where[Schema, View <: SchemaView]:
   val view: View
@@ -36,7 +36,7 @@ object Where:
     import quotes.reflect.*
 
     '{ ${ condition }(using ${ where }.view) } match
-      case '{ $cond: Col[BooleanOptType] } =>
+      case '{ $cond: Col[BooleanOptLike] } =>
         '{
           val filtered = ${ where }.underlying.where(${ cond }.untyped)
           StructDataFrame[Schema](filtered)

src/main/Where.scala

@@ -3,29 +3,46 @@ package api
 
 export DataFrameBuilders.toDF
 export types.{
-  DataType,
-  BooleanType,
-  BooleanOptType,
-  StringType,
-  StringOptType,
-  ByteType,
-  ByteOptType,
-  ShortType,
-  ShortOptType,
-  IntegerType,
-  IntegerOptType,
-  LongType,
-  LongOptType,
-  FloatType,
-  FloatOptType,
-  DoubleType,
-  DoubleOptType,
-  StructType,
-  StructOptType
+  boolean,
+  boolean_?,
+  BooleanNotNull,
+  BooleanOrNull,
+  string,
+  string_?,
+  StringNotNull,
+  StringOrNull,
+  byte,
+  byte_?,
+  ByteNotNull,
+  ByteOrNull,
+  short,
+  short_?,
+  ShortNotNull,
+  ShortOrNull,
+  int,
+  int_?,
+  IntNotNull,
+  IntOrNull,
+  long,
+  long_?,
+  LongNotNull,
+  LongOrNull,
+  float,
+  float_?,
+  FloatNotNull,
+  FloatOrNull,
+  double,
+  double_?,
+  DoubleNotNull,
+  DoubleOrNull,
+  struct,
+  struct_?,
+  StructNotNull,
+  StructOrNull
 }
 export UntypedOps.typed
 export org.virtuslab.iskra.$
-export org.virtuslab.iskra.{Column, Columns, DataFrame, ClassDataFrame, NamedColumns, StructDataFrame, UntypedColumn, UntypedDataFrame, :=, /}
+export org.virtuslab.iskra.{Column, Columns, Col, DataFrame, ClassDataFrame, NamedColumns, StructDataFrame, UntypedColumn, UntypedDataFrame, :=, /}
 
 object functions:
   export org.virtuslab.iskra.functions.{lit, when}

src/main/api/api.scala

@@ -5,22 +5,22 @@ import org.virtuslab.iskra.Agg
 import org.virtuslab.iskra.Col
 import org.virtuslab.iskra.UntypedOps.typed
 import org.virtuslab.iskra.types.*
-import org.virtuslab.iskra.types.DataType.{NumericOptType, Nullable}
+import org.virtuslab.iskra.types.DataType.AsNullable
 
 class Sum[A <: Agg](val agg: A):
-  def apply[T <: NumericOptType](column: agg.View ?=> Col[T]): Col[Nullable[T]] =
+  def apply[T <: DoubleOptLike](column: agg.View ?=> Col[T]): Col[AsNullable[T]] =
     sql.functions.sum(column(using agg.view).untyped).typed
 
 class Max[A <: Agg](val agg: A):
-  def apply[T <: NumericOptType](column: agg.View ?=> Col[T]): Col[Nullable[T]] =
+  def apply[T <: DoubleOptLike](column: agg.View ?=> Col[T]): Col[AsNullable[T]] =
     sql.functions.max(column(using agg.view).untyped).typed
 
 class Min[A <: Agg](val agg: A):
-  def apply[T <: NumericOptType](column: agg.View ?=> Col[T]): Col[Nullable[T]] =
+  def apply[T <: DoubleOptLike](column: agg.View ?=> Col[T]): Col[AsNullable[T]] =
     sql.functions.min(column(using agg.view).untyped).typed
 
 class Avg[A <: Agg](val agg: A):
-  def apply(column: agg.View ?=> Col[NumericOptType]): Col[DoubleOptType] =
+  def apply(column: agg.View ?=> Col[DoubleOptLike]): Col[DoubleOrNull] =
     sql.functions.avg(column(using agg.view).untyped).typed
 
 object Aggregates:

src/main/functions/aggregates.scala

@@ -1,17 +1,20 @@
 package org.virtuslab.iskra
 package types
 
-import DataType.{CommonNumericNonNullableType, CommonNumericNullableType, NumericOptType, NumericType}
-
-trait Coerce[-A <: DataType, -B <: DataType]:
+trait Coerce[A <: DataType, B <: DataType]:
   type Coerced <: DataType
 
-object Coerce:
-  given sameType[A <: DataType]: Coerce[A, A] with
+object Coerce extends CoerceLowPrio:
+  given sameType[A <: FinalDataType]: Coerce[A, A] with
+    override type Coerced = A
+
+  given nullableFirst[A <: FinalDataType & Nullable, B <: FinalDataType & NonNullable](using A <:< NullableOf[B]): Coerce[A, B] with
     override type Coerced = A
 
-  given nullable[A <: NumericOptType, B <: NumericOptType]: Coerce[A, B] with
-    override type Coerced = CommonNumericNullableType[A, B]
+  given nullableSecond[A <: FinalDataType & NonNullable, B <: FinalDataType & Nullable](using A <:< NonNullableOf[B]): Coerce[A, B] with
+    override type Coerced = B
 
-  given nonNullable[A <: NumericType, B <: NumericType]: Coerce[A, B] with
-    override type Coerced = CommonNumericNonNullableType[A, B]
+trait CoerceLowPrio:
+  given numeric[A <: FinalDataType & DoubleOptLike, B <: FinalDataType & DoubleOptLike]: (Coerce[A, B] { type Coerced = CommonNumericType[A, B] }) =
+    new Coerce[A, B]:
+      override type Coerced = CommonNumericType[A, B]