feat: migrate Expr to standalone enum

Desdaemon · Dec 30, 2023 · 0022ccc · 0022ccc
1 parent f2992a3
commit 0022ccc
Show file tree

Hide file tree

Showing 18 changed files with 20,501 additions and 11,027 deletions.
diff --git a/packages/polars/lib/src/expr.dart b/packages/polars/lib/src/expr.dart
@@ -5,71 +5,132 @@ final _kIsWeb = 0 == 0.0;
 
 extension LiteralValueExt on LiteralValue {
   /// Returns an expression representing this literal value.
-  Expr get expr => Expr.literal(value: this);
+  Expr get expr => Expr.literal(this);
 }
 
 /// Extensions for [Expr].
 extension ExprExt on Expr {
   /// Returns an expression evaluating whether this is less than [other].
-  Expr operator <(Expr other) => lt(other: other);
+  Expr operator <(Object? other) =>
+      Expr.binaryExpr(left: this, op: Operator.lt, right: other.expr);
 
   /// Returns an expression evaluating whether this is no greater than [other].
-  Expr operator <=(Expr other) => ltEq(other: other);
+  Expr operator <=(Object? other) =>
+      Expr.binaryExpr(left: this, op: Operator.ltEq, right: other.expr);
 
   /// Returns an expression evaluating whether this is greater than [other].
-  Expr operator >(Expr other) => gt(other: other);
+  Expr operator >(Object? other) =>
+      Expr.binaryExpr(left: this, op: Operator.gt, right: other.expr);
 
   /// Returns an expression evaluating whether this is no lesser than [other].
-  Expr operator >=(Expr other) => gtEq(other: other);
+  Expr operator >=(Object? other) =>
+      Expr.binaryExpr(left: this, op: Operator.gtEq, right: other.expr);
 
   /// Returns an expression representing the sum of this and [other].
-  Expr operator +(Expr other) => add(other: other);
+  Expr operator +(Object? other) =>
+      Expr.binaryExpr(left: this, op: Operator.plus, right: other.expr);
 
   /// Returns an expression representing the difference of this and [other].
-  Expr operator -(Expr other) => sub(other: other);
+  Expr operator -(Object? other) =>
+      Expr.binaryExpr(left: this, op: Operator.minus, right: other.expr);
 
   /// Returns an expression representing the product of this and [other].
-  Expr operator *(Expr other) => mul(other: other);
+  Expr operator *(Object? other) =>
+      Expr.binaryExpr(left: this, op: Operator.multiply, right: other.expr);
 
   /// Returns an expression representing the division of this and [other].
-  Expr operator /(Expr other) => div(other: other);
+  Expr operator /(Object? other) =>
+      Expr.binaryExpr(left: this, op: Operator.divide, right: other.expr);
 
   /// Returns an expression representing the integral division of this and [other].
-  Expr operator ~/(Expr other) => floorDiv(rhs: other);
+  Expr operator ~/(Object? other) =>
+      Expr.binaryExpr(left: this, op: Operator.floorDivide, right: other.expr);
 
   /// Performs a modulo operation on this and [other].
-  Expr operator %(Expr other) => rem(other: other);
+  Expr operator %(Object? other) =>
+      Expr.binaryExpr(left: this, op: Operator.modulus, right: other.expr);
 
   /// Returns an expression evaluating whether both this and [other] are true.
-  Expr operator &(Expr other) => and(expr: other);
+  Expr operator &(Object? other) =>
+      Expr.binaryExpr(left: this, op: Operator.and, right: other.expr);
 
   /// Returns an expression evaluating whether either this or [other] is true.
-  Expr operator |(Expr other) => or(expr: other);
+  Expr operator |(Object? other) =>
+      Expr.binaryExpr(left: this, op: Operator.or, right: other.expr);
 
   /// Returns an expression evaluating whether at most one of this and [other] is true,
   /// i.e. whether this and [other] are not equal.
-  Expr operator ^(Expr other) => xor(expr: other);
+  Expr operator ^(Object? other) =>
+      Expr.binaryExpr(left: this, op: Operator.xor, right: other.expr);
+
+  Expr head({int length = 10}) => Expr.slice(
+      input: this,
+      offset: const Expr.literal(LiteralValue.int64(0)),
+      length: length.expr);
+
+  Expr tail({int length = 10}) =>
+      Expr.slice(input: this, offset: (-length).expr, length: length.expr);
+
+  Expr alias(String name) => Expr.alias(this, name);
+  Expr get aggGroups => Expr.agg(AggExpr.aggGroups(this));
+  Expr cast(DataType dataType, {bool strict = false}) =>
+      Expr.cast(expr: this, dataType: dataType, strict: strict);
+  Expr equalMissing(Object? other) =>
+      Expr.binaryExpr(left: this, op: Operator.eqValidity, right: other.expr);
+  Expr notEqualMissing(Object? other) => Expr.binaryExpr(
+      left: this, op: Operator.notEqValidity, right: other.expr);
+  Expr exclude(Iterable<String> columns) =>
+      Expr.exclude(this, columns.map(Excluded.name).toList(growable: false));
+  Expr get explode => Expr.explode(this);
+
+  /// Alias for [explode].
+  Expr get flatten => Expr.explode(this);
+  Expr filter({required Object? by}) => Expr.filter(input: this, by: by.expr);
+  Expr get first => Expr.agg(AggExpr.first(this));
+  Expr get last => Expr.agg(AggExpr.last(this));
+  Expr get({required Object? index, bool returnsScalar = true}) =>
+      Expr.gather(expr: this, idx: index.expr, returnsScalar: returnsScalar);
+  Expr get implode => Expr.agg(AggExpr.implode(this));
+  Expr get nUnique => Expr.agg(AggExpr.nUnique(this));
+  Expr get nanMax => Expr.agg(AggExpr.max(input: this, propagateNans: true));
+  Expr get nanMin => Expr.agg(AggExpr.min(input: this, propagateNans: true));
+  Expr slice(int offset, int length) =>
+      Expr.slice(input: this, offset: offset.expr, length: length.expr);
+
+  /// Calculate the standard deviation of this expression with the specified
+  /// [dof] or [delta degrees of freedom](https://en.wikipedia.org/wiki/Degrees_of_freedom_(statistics)).
+  Expr std(int ddof) => Expr.agg(AggExpr.std(this, ddof));
 }
 
+Expr col(String column) => Expr.column(column);
+Expr cols(Iterable<String> columns) =>
+    Expr.columns(columns.toList(growable: false));
+
+Expr when(
+  Expr condition, {
+  required Object? then,
+  Object? otherwise = const Expr.literal(LiteralValue.Null()),
+}) =>
+    Expr.ternary(
+        predicate: condition, truthy: then.expr, falsy: otherwise.expr);
+
 /// Extensions on [String].
 extension StringPolars on String {
-  Expr get expr => Expr.literal(value: LiteralValue.utf8(this));
+  Expr get expr => Expr.literal(LiteralValue.utf8(this));
   DataType get dtype => const DataType.utf8();
 }
 
 /// Extensions on [int].
 extension IntPolars on int {
-  Expr get i32 => Expr.literal(value: LiteralValue.int32(this));
-  Expr get i64 => Expr.literal(value: LiteralValue.int64(this));
+  Expr get i32 => Expr.literal(LiteralValue.int32(this));
+  Expr get i64 => Expr.literal(LiteralValue.int64(this));
 
-  Expr get u32 =>
-      Expr.literal(value: LiteralValue.uint32(_assertNonNegative(this)));
-  Expr get u64 =>
-      Expr.literal(value: LiteralValue.uint64(_assertNonNegative(this)));
+  Expr get u32 => Expr.literal(LiteralValue.uint32(_assertNonNegative(this)));
+  Expr get u64 => Expr.literal(LiteralValue.uint64(_assertNonNegative(this)));
 
   Expr range(int other, {DataType? dataType}) {
     return Expr.literal(
-      value: LiteralValue.range(
+      LiteralValue.range(
         low: this,
         high: other >= this ? other : this,
         dataType: dataType ?? dtype,
@@ -83,51 +144,48 @@ extension IntPolars on int {
 
 /// Extensions on [double].
 extension DoublePolars on double {
-  Expr get f32 => Expr.literal(value: LiteralValue.float32(this));
-  Expr get f64 => Expr.literal(value: LiteralValue.float64(this));
+  Expr get f32 => Expr.literal(LiteralValue.float32(this));
+  Expr get f64 => Expr.literal(LiteralValue.float64(this));
   Expr get expr => f64;
   DataType get dtype => const DataType.float64();
 }
 
 /// Extensions on [bool].
 extension BoolPolars on bool {
-  Expr get expr => Expr.literal(
-      value: this
-          ? const LiteralValue.boolean(true)
-          : const LiteralValue.boolean(false));
+  Expr get expr => Expr.literal(this
+      ? const LiteralValue.boolean(true)
+      : const LiteralValue.boolean(false));
   DataType get dtype => const DataType.boolean();
 }
 
 /// Extensions on [DateTime].
 extension DateTimePolars on DateTime {
   DataType get dtype => DataType.datetime(
       _kIsWeb ? TimeUnit.milliseconds : TimeUnit.microseconds);
-  Expr get expr => Expr.literal(
-      value: _kIsWeb
-          ? LiteralValue.dateTime(
-              millisecondsSinceEpoch,
-              TimeUnit.milliseconds,
-            )
-          : LiteralValue.dateTime(
-              microsecondsSinceEpoch,
-              TimeUnit.microseconds,
-            ));
+  Expr get expr => Expr.literal(_kIsWeb
+      ? LiteralValue.dateTime(
+          millisecondsSinceEpoch,
+          TimeUnit.milliseconds,
+        )
+      : LiteralValue.dateTime(
+          microsecondsSinceEpoch,
+          TimeUnit.microseconds,
+        ));
 }
 
 /// Extensions on [Duration].
 extension DurationPolars on Duration {
   DataType get dtype => DataType.duration(
       _kIsWeb ? TimeUnit.milliseconds : TimeUnit.microseconds);
-  Expr get expr => Expr.literal(
-      value: _kIsWeb
-          ? LiteralValue.duration(
-              inMilliseconds,
-              TimeUnit.milliseconds,
-            )
-          : LiteralValue.duration(
-              inMicroseconds,
-              TimeUnit.microseconds,
-            ));
+  Expr get expr => Expr.literal(_kIsWeb
+      ? LiteralValue.duration(
+          inMilliseconds,
+          TimeUnit.milliseconds,
+        )
+      : LiteralValue.duration(
+          inMicroseconds,
+          TimeUnit.microseconds,
+        ));
 }
 
 extension DynamicPolars on dynamic {
@@ -138,7 +196,9 @@ extension DynamicPolars on dynamic {
         bool value => value.expr,
         DateTime value => value.expr,
         Duration value => value.expr,
-        null => Expr.literal(value: const LiteralValue.Null()),
+        Expr expr => expr,
+        LiteralValue lit => Expr.literal(lit),
+        null => const Expr.literal(LiteralValue.Null()),
         _ => '$this'.expr,
       };
   DataType get dtype => switch (this) {