Skip to content

Commit 47f7519

Browse files
mchavmchav
committed
chore: Add foldlColumn and foldl1Column.
1 parent f12bd45 commit 47f7519

File tree

2 files changed

+89
-10
lines changed

2 files changed

+89
-10
lines changed

‎src/DataFrame/Internal/Column.hs‎

Lines changed: 82 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -606,6 +606,88 @@ ifoldlColumn f acc c@(UnboxedColumn (column :: VU.Vector d)) = case testEquality
606606
}
607607
)
608608

609+
foldlColumn ::
610+
forall a b.
611+
(Columnable a, Columnable b) =>
612+
(b -> a -> b) -> b -> Column -> Either DataFrameException b
613+
foldlColumn f acc c@(BoxedColumn (column :: VB.Vector d)) = case testEquality (typeRep @a) (typeRep @d) of
614+
Just Refl -> pure $ VG.foldl' f acc column
615+
Nothing ->
616+
Left $
617+
TypeMismatchException
618+
( MkTypeErrorContext
619+
{ userType = Right (typeRep @a)
620+
, expectedType = Right (typeRep @d)
621+
, callingFunctionName = Just "foldlColumn"
622+
, errorColumnName = Nothing
623+
}
624+
)
625+
foldlColumn f acc c@(OptionalColumn (column :: VB.Vector d)) = case testEquality (typeRep @a) (typeRep @d) of
626+
Just Refl -> pure $ VG.foldl' f acc column
627+
Nothing ->
628+
Left $
629+
TypeMismatchException
630+
( MkTypeErrorContext
631+
{ userType = Right (typeRep @a)
632+
, expectedType = Right (typeRep @d)
633+
, callingFunctionName = Just "foldlColumn"
634+
, errorColumnName = Nothing
635+
}
636+
)
637+
foldlColumn f acc c@(UnboxedColumn (column :: VU.Vector d)) = case testEquality (typeRep @a) (typeRep @d) of
638+
Just Refl -> pure $ VG.foldl' f acc column
639+
Nothing ->
640+
Left $
641+
TypeMismatchException
642+
( MkTypeErrorContext
643+
{ userType = Right (typeRep @a)
644+
, expectedType = Right (typeRep @d)
645+
, callingFunctionName = Just "foldlColumn"
646+
, errorColumnName = Nothing
647+
}
648+
)
649+
650+
foldl1Column ::
651+
forall a.
652+
(Columnable a) =>
653+
(a -> a -> a) -> Column -> Either DataFrameException a
654+
foldl1Column f c@(BoxedColumn (column :: VB.Vector d)) = case testEquality (typeRep @a) (typeRep @d) of
655+
Just Refl -> pure $ VG.foldl1' f column
656+
Nothing ->
657+
Left $
658+
TypeMismatchException
659+
( MkTypeErrorContext
660+
{ userType = Right (typeRep @a)
661+
, expectedType = Right (typeRep @d)
662+
, callingFunctionName = Just "foldl1Column"
663+
, errorColumnName = Nothing
664+
}
665+
)
666+
foldl1Column f c@(OptionalColumn (column :: VB.Vector d)) = case testEquality (typeRep @a) (typeRep @d) of
667+
Just Refl -> pure $ VG.foldl1' f column
668+
Nothing ->
669+
Left $
670+
TypeMismatchException
671+
( MkTypeErrorContext
672+
{ userType = Right (typeRep @a)
673+
, expectedType = Right (typeRep @d)
674+
, callingFunctionName = Just "foldl1Column"
675+
, errorColumnName = Nothing
676+
}
677+
)
678+
foldl1Column f c@(UnboxedColumn (column :: VU.Vector d)) = case testEquality (typeRep @a) (typeRep @d) of
679+
Just Refl -> pure $ VG.foldl1' f column
680+
Nothing ->
681+
Left $
682+
TypeMismatchException
683+
( MkTypeErrorContext
684+
{ userType = Right (typeRep @a)
685+
, expectedType = Right (typeRep @d)
686+
, callingFunctionName = Just "foldl1Column"
687+
, errorColumnName = Nothing
688+
}
689+
)
690+
609691
headColumn :: forall a. (Columnable a) => Column -> Either DataFrameException a
610692
headColumn (BoxedColumn (col :: VB.Vector b)) = case testEquality (typeRep @a) (typeRep @b) of
611693
Just Refl ->

‎src/DataFrame/Internal/Expression.hs‎

Lines changed: 7 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -123,8 +123,7 @@ interpret df expression@(AggVector expr op (f :: v b -> c)) = do
123123
(UnboxedColumn col) -> processColumn col
124124
interpret df expression@(AggReduce expr op (f :: a -> a -> a)) = first (handleInterpretException (show expr)) $ do
125125
(TColumn column) <- interpret @a df expr
126-
h <- headColumn @a column
127-
value <- ifoldlColumn (\acc _ v -> f acc v) h column
126+
value <- foldl1Column f column
128127
pure $ TColumn $ fromVector $ V.replicate (fst $ dataframeDimensions df) value
129128
interpret df expression@(AggNumericVector expr op (f :: VU.Vector b -> c)) = first (handleInterpretException (show expression)) $ do
130129
(TColumn column) <- interpret @b df expr
@@ -144,7 +143,7 @@ interpret df expression@(AggNumericVector expr op (f :: VU.Vector b -> c)) = fir
144143
_ -> error "Trying to apply numeric computation to non-numeric column"
145144
interpret df expression@(AggFold expr op start (f :: (a -> b -> a))) = first (handleInterpretException (show expression)) $ do
146145
(TColumn column) <- interpret @b df expr
147-
value <- ifoldlColumn (\acc _ v -> f acc v) start column
146+
value <- foldlColumn f start column
148147
pure $ TColumn $ fromVector $ V.replicate (fst $ dataframeDimensions df) value
149148

150149
data AggregationResult a
@@ -793,20 +792,18 @@ interpretAggregation gdf@(Grouped df names indices os) expression@(AggReduce exp
793792
Aggregated $
794793
TColumn $
795794
fromVector $
796-
V.map (\v -> VU.foldl' f (VG.head v) (VG.drop 1 v)) col
795+
V.map (VU.foldl1' f) col
797796
SFalse -> Left $ InternalException "Boxed type inside an unboxed column"
798797
Just Refl ->
799798
Right $
800799
Aggregated $
801800
TColumn $
802801
fromVector $
803-
V.map (\v -> VG.foldl' f (VG.head v) (VG.drop 1 v)) col
802+
V.map (VG.foldl1' f) col
804803
Right (UnAggregated _) -> Left $ InternalException "Aggregated into non-boxed column"
805-
Right (Aggregated (TColumn column)) -> case headColumn @a column of
804+
Right (Aggregated (TColumn column)) -> case foldl1Column f column of
806805
Left e -> Left e
807-
Right h -> case ifoldlColumn (\acc _ v -> f acc v) h column of
808-
Left e -> Left e
809-
Right value -> interpretAggregation @a gdf (Lit value)
806+
Right value -> interpretAggregation @a gdf (Lit value)
810807
interpretAggregation gdf@(Grouped df names indices os) expression@(AggFold expr op s (f :: (a -> b -> a))) =
811808
case interpretAggregation @b gdf expr of
812809
(Left (TypeMismatchException context)) ->
@@ -827,7 +824,7 @@ interpretAggregation gdf@(Grouped df names indices os) expression@(AggFold expr
827824
SFalse -> Left $ InternalException "Boxed type inside an unboxed column"
828825
Nothing -> Left $ nestedTypeException @d @b (show expr)
829826
Right (UnAggregated _) -> Left $ InternalException "Aggregated into non-boxed column"
830-
Right (Aggregated (TColumn column)) -> case ifoldlColumn (\acc _ v -> f acc v) s column of
827+
Right (Aggregated (TColumn column)) -> case foldlColumn f s column of
831828
Left e -> Left e
832829
Right value -> interpretAggregation @a gdf (Lit value)
833830

0 commit comments

Comments
 (0)