11package org .virtuslab .iskra
22package types
33
4- import scala .quoted ._
5- import scala .deriving .Mirror
6- import org .apache .spark .sql
7- import MacroHelpers .TupleSubtype
8-
94sealed trait DataType
105
116object DataType :
@@ -38,7 +33,7 @@ object DataType:
3833 case DoubleOptType => DoubleOptType
3934 case StructOptType [schema] => StructOptType [schema]
4035
41- type CommonNumericOptType [T1 <: DataType , T2 <: DataType ] <: NumericOptType = (T1 , T2 ) match
36+ type CommonNumericNullableType [T1 <: DataType , T2 <: DataType ] <: NumericOptType = (T1 , T2 ) match
4237 case (DoubleOptType , _) | (_, DoubleOptType ) => DoubleOptType
4338 case (FloatOptType , _) | (_, FloatOptType ) => FloatOptType
4439 case (LongOptType , _) | (_, LongOptType ) => LongOptType
@@ -54,179 +49,6 @@ object DataType:
5449 case (ShortOptType , _) | (_, ShortOptType ) => ShortType
5550 case (ByteOptType , _) | (_, ByteOptType ) => ByteType
5651
57- trait Encoder [- A ]:
58- type ColumnType <: DataType
59- def encode (value : A ): Any
60- def decode (value : Any ): Any
61- def catalystType : sql.types.DataType
62- def isNullable : Boolean
63-
64- trait PrimitiveEncoder [- A ] extends Encoder [A ]
65-
66- trait PrimitiveNullableEncoder [- A ] extends PrimitiveEncoder [Option [A ]]:
67- def encode (value : Option [A ]) = value.orNull
68- def decode (value : Any ) = Option (value)
69- def isNullable = true
70-
71- trait PrimitiveNonNullableEncoder [- A ] extends PrimitiveEncoder [A ]:
72- def encode (value : A ) = value
73- def decode (value : Any ) = value
74- def isNullable = true
75-
76-
77- object Encoder :
78- type Aux [- A , E <: DataType ] = Encoder [A ] { type ColumnType = E }
79-
80- inline given boolean : PrimitiveNonNullableEncoder [Boolean ] with
81- type ColumnType = BooleanType
82- def catalystType = sql.types.BooleanType
83- inline given booleanOpt : PrimitiveNullableEncoder [Boolean ] with
84- type ColumnType = BooleanOptType
85- def catalystType = sql.types.BooleanType
86-
87- inline given string : PrimitiveNonNullableEncoder [String ] with
88- type ColumnType = StringType
89- def catalystType = sql.types.StringType
90- inline given stringOpt : PrimitiveNullableEncoder [String ] with
91- type ColumnType = StringOptType
92- def catalystType = sql.types.StringType
93-
94- inline given byte : PrimitiveNonNullableEncoder [Byte ] with
95- type ColumnType = ByteType
96- def catalystType = sql.types.ByteType
97- inline given byteOpt : PrimitiveNullableEncoder [Byte ] with
98- type ColumnType = ByteOptType
99- def catalystType = sql.types.ByteType
100-
101- inline given short : PrimitiveNonNullableEncoder [Short ] with
102- type ColumnType = ShortType
103- def catalystType = sql.types.ShortType
104- inline given shortOpt : PrimitiveNullableEncoder [Short ] with
105- type ColumnType = ShortOptType
106- def catalystType = sql.types.ShortType
107-
108- inline given int : PrimitiveNonNullableEncoder [Int ] with
109- type ColumnType = IntegerType
110- def catalystType = sql.types.IntegerType
111- inline given intOpt : PrimitiveNullableEncoder [Int ] with
112- type ColumnType = IntegerOptType
113- def catalystType = sql.types.IntegerType
114-
115- inline given long : PrimitiveNonNullableEncoder [Long ] with
116- type ColumnType = LongType
117- def catalystType = sql.types.LongType
118- inline given longOpt : PrimitiveNullableEncoder [Long ] with
119- type ColumnType = LongOptType
120- def catalystType = sql.types.LongType
121-
122- inline given float : PrimitiveNonNullableEncoder [Float ] with
123- type ColumnType = FloatType
124- def catalystType = sql.types.FloatType
125- inline given floatOpt : PrimitiveNullableEncoder [Float ] with
126- type ColumnType = FloatOptType
127- def catalystType = sql.types.FloatType
128-
129- inline given double : PrimitiveNonNullableEncoder [Double ] with
130- type ColumnType = DoubleType
131- def catalystType = sql.types.DoubleType
132- inline given doubleOpt : PrimitiveNullableEncoder [Double ] with
133- type ColumnType = DoubleOptType
134- def catalystType = sql.types.DoubleType
135-
136- export StructEncoder .{fromMirror , optFromMirror }
137-
138- trait StructEncoder [- A ] extends Encoder [A ]:
139- type StructSchema <: Tuple
140- type ColumnType = StructType [StructSchema ]
141- override def catalystType : sql.types.StructType
142- override def encode (a : A ): sql.Row
143-
144- object StructEncoder :
145- type Aux [- A , E <: Tuple ] = StructEncoder [A ] { type StructSchema = E }
146-
147- private case class ColumnInfo (
148- labelType : Type [? ],
149- labelValue : String ,
150- decodedType : Type [? ],
151- encoder : Expr [Encoder [? ]]
152- )
153-
154- private def getColumnSchemaType (using quotes : Quotes )(subcolumnInfos : List [ColumnInfo ]): Type [? ] =
155- subcolumnInfos match
156- case Nil => Type .of[EmptyTuple ]
157- case info :: tail =>
158- info.labelType match
159- case ' [Name .Subtype [lt]] =>
160- info.encoder match
161- case ' { $ {encoder}: Encoder .Aux [tpe, DataType .Subtype [e]] } =>
162- getColumnSchemaType(tail) match
163- case ' [TupleSubtype [tailType]] =>
164- Type .of[(lt := e) *: tailType]
165-
166- private def getSubcolumnInfos (elemLabels : Type [? ], elemTypes : Type [? ])(using Quotes ): List [ColumnInfo ] =
167- import quotes .reflect .Select
168- elemLabels match
169- case ' [EmptyTuple ] => Nil
170- case ' [label *: labels] =>
171- val labelValue = Type .valueOfConstant[label].get.toString
172- elemTypes match
173- case ' [tpe *: types] =>
174- Expr .summon[Encoder [tpe]] match
175- case Some (encoderExpr) =>
176- ColumnInfo (Type .of[label], labelValue, Type .of[tpe], encoderExpr) :: getSubcolumnInfos(Type .of[labels], Type .of[types])
177- case _ => quotes.reflect.report.errorAndAbort(s " Could not summon encoder for ${Type .show[tpe]}" )
178-
179- transparent inline given fromMirror [A ]: StructEncoder [A ] = $ { fromMirrorImpl[A ] }
180-
181- def fromMirrorImpl [A : Type ](using q : Quotes ): Expr [StructEncoder [A ]] =
182- Expr .summon[Mirror .Of [A ]].getOrElse(throw new Exception (s " Could not find Mirror when generating encoder for ${Type .show[A ]}" )) match
183- case ' { $ {mirror}: Mirror .ProductOf [A ] { type MirroredElemLabels = elementLabels; type MirroredElemTypes = elementTypes } } =>
184- val subcolumnInfos = getSubcolumnInfos(Type .of[elementLabels], Type .of[elementTypes])
185- val columnSchemaType = getColumnSchemaType(subcolumnInfos)
186-
187- val structFieldExprs = subcolumnInfos.map { info =>
188- ' { sql.types.StructField ($ {Expr (info.labelValue)}, $ {info.encoder}.catalystType, $ {info.encoder}.isNullable) }
189- }
190- val structFields = Expr .ofSeq(structFieldExprs)
191-
192- def rowElements (value : Expr [A ]) =
193- subcolumnInfos.map { info =>
194- import quotes .reflect .*
195- info.decodedType match
196- case ' [t] =>
197- ' { $ {info.encoder}.asInstanceOf [Encoder [t]].encode($ { Select .unique(value.asTerm, info.labelValue).asExprOf[t] }) }
198- }
199-
200- def rowElementsTuple (row : Expr [sql.Row ]): Expr [Tuple ] =
201- val elements = subcolumnInfos.zipWithIndex.map { (info, idx) =>
202- given Quotes = q
203- ' { $ {info.encoder}.decode($ {row}.get($ {Expr (idx)})) }
204- }
205- Expr .ofTupleFromSeq(elements)
206-
207- columnSchemaType match
208- case ' [TupleSubtype [t]] =>
209- ' {
210- (new StructEncoder [A ] {
211- override type StructSchema = t
212- override def catalystType = sql.types.StructType ($ { structFields })
213- override def isNullable = false
214- override def encode (a : A ) =
215- sql.Row .fromSeq($ { Expr .ofSeq(rowElements(' a )) })
216- override def decode (a : Any ) =
217- $ {mirror}.fromProduct($ { rowElementsTuple(' {a.asInstanceOf [sql.Row ]}) })
218- }): StructEncoder [A ] { type StructSchema = t }
219- }
220- end fromMirrorImpl
221-
222- inline given optFromMirror [A ](using encoder : StructEncoder [A ]): (Encoder [Option [A ]] { type ColumnType = StructOptType [encoder.StructSchema ] }) =
223- new Encoder [Option [A ]]:
224- override type ColumnType = StructOptType [encoder.StructSchema ]
225- override def encode (value : Option [A ]): Any = value.map(encoder.encode).orNull
226- override def decode (value : Any ): Any = Option (encoder.decode)
227- override def catalystType = encoder.catalystType
228- override def isNullable = true
229-
23052import DataType .NotNull
23153
23254sealed class BooleanOptType extends DataType
0 commit comments