Skip to content

Commit

Permalink
Make actual data dypes disjunct
Browse files Browse the repository at this point in the history
  • Loading branch information
prolativ committed Oct 23, 2024
1 parent d3a6c44 commit c047d70
Show file tree
Hide file tree
Showing 16 changed files with 380 additions and 382 deletions.
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
2 changes: 1 addition & 1 deletion src/main/FrameSchema.scala
Original file line number Diff line number Diff line change
Expand Up @@ -25,7 +25,7 @@ object FrameSchema:
(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 Tuple => Tuple.Map[T, NullableLabeledDataType]
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: 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

0 comments on commit c047d70

Please sign in to comment.