Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Better match types #39

Merged
merged 2 commits into from
Oct 23, 2024
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
307 changes: 114 additions & 193 deletions src/main/ColumnOp.scala

Large diffs are not rendered by default.

2 changes: 1 addition & 1 deletion src/main/DataFrameBuilders.scala
Original file line number Diff line number Diff line change
Expand Up @@ -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])
Expand Down
30 changes: 20 additions & 10 deletions src/main/FrameSchema.scala
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
4 changes: 2 additions & 2 deletions src/main/JoinOnCondition.scala
Original file line number Diff line number Diff line change
Expand Up @@ -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]

Expand Down Expand Up @@ -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)
Expand Down
4 changes: 0 additions & 4 deletions src/main/MacroHelpers.scala
Original file line number Diff line number Diff line change
Expand Up @@ -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
4 changes: 2 additions & 2 deletions src/main/SchemaView.scala
Original file line number Diff line number Diff line change
Expand Up @@ -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

Expand Down Expand Up @@ -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))
Expand Down
8 changes: 1 addition & 7 deletions src/main/Select.scala
Original file line number Diff line number Diff line change
Expand Up @@ -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))
Expand Down
9 changes: 8 additions & 1 deletion src/main/StructDataFrame.scala
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -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 _ =>
Expand Down
4 changes: 2 additions & 2 deletions src/main/UntypedOps.scala
Original file line number Diff line number Diff line change
@@ -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)
Expand All @@ -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 }) }
8 changes: 4 additions & 4 deletions src/main/When.scala
Original file line number Diff line number Diff line change
@@ -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))
4 changes: 2 additions & 2 deletions src/main/Where.scala
Original file line number Diff line number Diff line change
@@ -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
Expand Down Expand Up @@ -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)
Expand Down
57 changes: 37 additions & 20 deletions src/main/api/api.scala
Original file line number Diff line number Diff line change
Expand Up @@ -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}
Expand Down
10 changes: 5 additions & 5 deletions src/main/functions/aggregates.scala
Original file line number Diff line number Diff line change
Expand Up @@ -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:
Expand Down
21 changes: 12 additions & 9 deletions src/main/types/Coerce.scala
Original file line number Diff line number Diff line change
@@ -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]
Loading
Loading