|
1 | 1 | package org.virtuslab.iskra
|
2 | 2 |
|
| 3 | +import scala.language.implicitConversions |
| 4 | + |
3 | 5 | import scala.quoted.*
|
4 | 6 |
|
5 | 7 | import org.apache.spark.sql.{Column => UntypedColumn}
|
6 | 8 | import types.DataType
|
| 9 | +import MacroHelpers.TupleSubtype |
| 10 | + |
| 11 | +class Column(val untyped: UntypedColumn): |
| 12 | + inline def name(using v: ValueOf[Name]): Name = v.value |
| 13 | + |
| 14 | +object Column: |
| 15 | + implicit transparent inline def columnToNamedColumn(inline col: Col[?]): NamedColumn[?, ?] = |
| 16 | + ${ columnToNamedColumnImpl('col) } |
| 17 | + |
| 18 | + private def columnToNamedColumnImpl(col: Expr[Col[?]])(using Quotes): Expr[NamedColumn[?, ?]] = |
| 19 | + import quotes.reflect.* |
| 20 | + col match |
| 21 | + case '{ ($v: StructuralSchemaView).selectDynamic($nm: Name).$asInstanceOf$[Col[tp]] } => |
| 22 | + nm.asTerm.tpe.asType match |
| 23 | + case '[Name.Subtype[n]] => |
| 24 | + '{ NamedColumn[n, tp](${ col }.untyped.as(${ nm })) } |
| 25 | + case '{ $c: Col[tp] } => |
| 26 | + col.asTerm match |
| 27 | + case Inlined(_, _, Ident(name)) => |
| 28 | + ConstantType(StringConstant(name)).asType match |
| 29 | + case '[Name.Subtype[n]] => |
| 30 | + val alias = Literal(StringConstant(name)).asExprOf[Name] |
| 31 | + '{ NamedColumn[n, tp](${ col }.untyped.as(${ alias })) } |
| 32 | + |
| 33 | + extension [T <: DataType](col: Col[T]) |
| 34 | + inline def as[N <: Name](name: N): NamedColumn[N, T] = |
| 35 | + NamedColumn[N, T](col.untyped.as(name)) |
| 36 | + inline def alias[N <: Name](name: N): NamedColumn[N, T] = |
| 37 | + NamedColumn[N, T](col.untyped.as(name)) |
| 38 | + |
| 39 | + extension [T1 <: DataType](col1: Col[T1]) |
| 40 | + inline def +[T2 <: DataType](col2: Col[T2])(using op: ColumnOp.Plus[T1, T2]): Col[op.Out] = op(col1, col2) |
| 41 | + inline def -[T2 <: DataType](col2: Col[T2])(using op: ColumnOp.Minus[T1, T2]): Col[op.Out] = op(col1, col2) |
| 42 | + inline def *[T2 <: DataType](col2: Col[T2])(using op: ColumnOp.Mult[T1, T2]): Col[op.Out] = op(col1, col2) |
| 43 | + inline def /[T2 <: DataType](col2: Col[T2])(using op: ColumnOp.Div[T1, T2]): Col[op.Out] = op(col1, col2) |
| 44 | + inline def ++[T2 <: DataType](col2: Col[T2])(using op: ColumnOp.PlusPlus[T1, T2]): Col[op.Out] = op(col1, col2) |
| 45 | + inline def <[T2 <: DataType](col2: Col[T2])(using op: ColumnOp.Lt[T1, T2]): Col[op.Out] = op(col1, col2) |
| 46 | + inline def <=[T2 <: DataType](col2: Col[T2])(using op: ColumnOp.Le[T1, T2]): Col[op.Out] = op(col1, col2) |
| 47 | + inline def >[T2 <: DataType](col2: Col[T2])(using op: ColumnOp.Gt[T1, T2]): Col[op.Out] = op(col1, col2) |
| 48 | + inline def >=[T2 <: DataType](col2: Col[T2])(using op: ColumnOp.Ge[T1, T2]): Col[op.Out] = op(col1, col2) |
| 49 | + inline def ===[T2 <: DataType](col2: Col[T2])(using op: ColumnOp.Eq[T1, T2]): Col[op.Out] = op(col1, col2) |
| 50 | + inline def =!=[T2 <: DataType](col2: Col[T2])(using op: ColumnOp.Ne[T1, T2]): Col[op.Out] = op(col1, col2) |
| 51 | + inline def &&[T2 <: DataType](col2: Col[T2])(using op: ColumnOp.And[T1, T2]): Col[op.Out] = op(col1, col2) |
| 52 | + inline def ||[T2 <: DataType](col2: Col[T2])(using op: ColumnOp.Or[T1, T2]): Col[op.Out] = op(col1, col2) |
| 53 | + |
| 54 | + |
| 55 | +class Col[+T <: DataType](untyped: UntypedColumn) extends Column(untyped) |
7 | 56 |
|
8 |
| -sealed trait NamedColumns[Schema](val underlyingColumns: Seq[UntypedColumn]) |
9 | 57 |
|
10 | 58 | object Columns:
|
11 |
| - transparent inline def apply(inline columns: NamedColumns[?]*): NamedColumns[?] = ${ applyImpl('columns) } |
| 59 | + transparent inline def apply[C <: NamedColumns](columns: C): ColumnsWithSchema[?] = ${ applyImpl('columns) } |
12 | 60 |
|
13 |
| - private def applyImpl(columns: Expr[Seq[NamedColumns[?]]])(using Quotes): Expr[NamedColumns[?]] = |
| 61 | + private def applyImpl[C : Type](columns: Expr[C])(using Quotes): Expr[ColumnsWithSchema[?]] = |
14 | 62 | import quotes.reflect.*
|
15 | 63 |
|
16 |
| - val columnValuesWithTypes = columns match |
17 |
| - case Varargs(colExprs) => |
18 |
| - colExprs.map { arg => |
19 |
| - arg match |
20 |
| - case '{ $value: NamedColumns[schema] } => ('{ ${ value }.underlyingColumns }, Type.of[schema]) |
21 |
| - } |
| 64 | + Expr.summon[CollectColumns[C]] match |
| 65 | + case Some(collectColumns) => |
| 66 | + collectColumns match |
| 67 | + case '{ $cc: CollectColumns[?] { type CollectedColumns = collectedColumns } } => |
| 68 | + Type.of[collectedColumns] match |
| 69 | + case '[TupleSubtype[collectedCols]] => |
| 70 | + '{ |
| 71 | + val cols = ${ cc }.underlyingColumns(${ columns }) |
| 72 | + ColumnsWithSchema[collectedCols](cols) |
| 73 | + } |
| 74 | + case None => |
| 75 | + throw CollectColumns.CannotCollectColumns(Type.show[C]) |
22 | 76 |
|
23 |
| - val columnsValues = columnValuesWithTypes.map(_._1) |
24 |
| - val columnsTypes = columnValuesWithTypes.map(_._2) |
25 | 77 |
|
26 |
| - val schemaTpe = FrameSchema.schemaTypeFromColumnsTypes(columnsTypes) |
| 78 | +trait NamedColumnOrColumnsLike |
27 | 79 |
|
28 |
| - schemaTpe match |
29 |
| - case '[s] => |
30 |
| - '{ |
31 |
| - val cols = ${ Expr.ofSeq(columnsValues) }.flatten |
32 |
| - new NamedColumns[s](cols) {} |
33 |
| - } |
| 80 | +type NamedColumns = Repeated[NamedColumnOrColumnsLike] |
34 | 81 |
|
35 |
| -class Column[+T <: DataType](val untyped: UntypedColumn): |
| 82 | +class NamedColumn[N <: Name, T <: DataType](val untyped: UntypedColumn) |
| 83 | + extends NamedColumnOrColumnsLike |
36 | 84 |
|
37 |
| - inline def name(using v: ValueOf[Name]): Name = v.value |
| 85 | +class ColumnsWithSchema[Schema <: Tuple](val underlyingColumns: Seq[UntypedColumn]) extends NamedColumnOrColumnsLike |
38 | 86 |
|
39 |
| -object Column: |
40 |
| - extension [T <: DataType](col: Column[T]) |
41 |
| - inline def as[N <: Name](name: N)(using v: ValueOf[N]): LabeledColumn[N, T] = |
42 |
| - LabeledColumn[N, T](col.untyped.as(v.value)) |
43 |
| - inline def alias[N <: Name](name: N)(using v: ValueOf[N]): LabeledColumn[N, T] = |
44 |
| - LabeledColumn[N, T](col.untyped.as(v.value)) |
45 |
| - |
46 |
| - extension [T1 <: DataType](col1: Column[T1]) |
47 |
| - inline def +[T2 <: DataType](col2: Column[T2])(using op: ColumnOp.Plus[T1, T2]): Column[op.Out] = op(col1, col2) |
48 |
| - inline def -[T2 <: DataType](col2: Column[T2])(using op: ColumnOp.Minus[T1, T2]): Column[op.Out] = op(col1, col2) |
49 |
| - inline def *[T2 <: DataType](col2: Column[T2])(using op: ColumnOp.Mult[T1, T2]): Column[op.Out] = op(col1, col2) |
50 |
| - inline def /[T2 <: DataType](col2: Column[T2])(using op: ColumnOp.Div[T1, T2]): Column[op.Out] = op(col1, col2) |
51 |
| - inline def ++[T2 <: DataType](col2: Column[T2])(using op: ColumnOp.PlusPlus[T1, T2]): Column[op.Out] = op(col1, col2) |
52 |
| - inline def <[T2 <: DataType](col2: Column[T2])(using op: ColumnOp.Lt[T1, T2]): Column[op.Out] = op(col1, col2) |
53 |
| - inline def <=[T2 <: DataType](col2: Column[T2])(using op: ColumnOp.Le[T1, T2]): Column[op.Out] = op(col1, col2) |
54 |
| - inline def >[T2 <: DataType](col2: Column[T2])(using op: ColumnOp.Gt[T1, T2]): Column[op.Out] = op(col1, col2) |
55 |
| - inline def >=[T2 <: DataType](col2: Column[T2])(using op: ColumnOp.Ge[T1, T2]): Column[op.Out] = op(col1, col2) |
56 |
| - inline def ===[T2 <: DataType](col2: Column[T2])(using op: ColumnOp.Eq[T1, T2]): Column[op.Out] = op(col1, col2) |
57 |
| - inline def =!=[T2 <: DataType](col2: Column[T2])(using op: ColumnOp.Ne[T1, T2]): Column[op.Out] = op(col1, col2) |
58 |
| - inline def &&[T2 <: DataType](col2: Column[T2])(using op: ColumnOp.And[T1, T2]): Column[op.Out] = op(col1, col2) |
59 |
| - inline def ||[T2 <: DataType](col2: Column[T2])(using op: ColumnOp.Or[T1, T2]): Column[op.Out] = op(col1, col2) |
60 | 87 |
|
61 | 88 | @annotation.showAsInfix
|
62 |
| -class :=[L <: LabeledColumn.Label, T <: DataType](untyped: UntypedColumn) |
63 |
| - extends Column[T](untyped) |
64 |
| - with NamedColumns[(L := T) *: EmptyTuple](Seq(untyped)) |
| 89 | +trait :=[L <: ColumnLabel, T <: DataType] |
65 | 90 |
|
66 | 91 | @annotation.showAsInfix
|
67 | 92 | trait /[+Prefix <: Name, +Suffix <: Name]
|
68 | 93 |
|
69 |
| -type LabeledColumn[L <: LabeledColumn.Label, T <: DataType] = :=[L, T] |
70 |
| - |
71 |
| -object LabeledColumn: |
72 |
| - type Label = Name | (Name / Name) |
73 |
| - def apply[L <: LabeledColumn.Label, T <: DataType](untyped: UntypedColumn) = new :=[L, T](untyped) |
| 94 | +type ColumnLabel = Name | (Name / Name) |
0 commit comments