diff --git a/core/src/main/kotlin/org/jetbrains/kotlinx/dataframe/api/with.kt b/core/src/main/kotlin/org/jetbrains/kotlinx/dataframe/api/with.kt index 190d0cfe7e..406947ab1c 100644 --- a/core/src/main/kotlin/org/jetbrains/kotlinx/dataframe/api/with.kt +++ b/core/src/main/kotlin/org/jetbrains/kotlinx/dataframe/api/with.kt @@ -4,13 +4,60 @@ import org.jetbrains.kotlinx.dataframe.AnyColumnReference import org.jetbrains.kotlinx.dataframe.DataFrame import org.jetbrains.kotlinx.dataframe.DataRow import org.jetbrains.kotlinx.dataframe.RowExpression +import org.jetbrains.kotlinx.dataframe.documentation.DocumentationUrls +import org.jetbrains.kotlinx.dataframe.documentation.ExcludeFromSources import org.jetbrains.kotlinx.dataframe.impl.aggregation.internal import org.jetbrains.kotlinx.dataframe.impl.aggregation.withExpr import org.jetbrains.kotlinx.dataframe.impl.emptyPath import kotlin.reflect.typeOf +// region docs + +/** + * For each pivot cell: an empty group yields `null` (type `V?`), a group of exactly one row yields + * that single value (type `V`), and a group of several rows yields a `List`. + */ +@ExcludeFromSources +private typealias PivotWithOneOrManyResultSnippet = Nothing + +/** + * For each pivot cell: a group with no row selected by the reducer yields `null` (type `V?`), + * otherwise the value produced by the expression (type `V`). + */ +@ExcludeFromSources +private typealias PivotWithReducedResultSnippet = Nothing + +/** + * If [expression\] returns a column reference, it is resolved to that column's value in the row + * the expression was evaluated on. + */ +@ExcludeFromSources +private typealias PivotWithColumnReferenceSnippet = Nothing + +// endregion + // region Pivot +/** + * Aggregates this [Pivot] by evaluating the given [row expression][expression] on every row of each pivot group. + * + * @include [PivotWithOneOrManyResultSnippet] + * + * @include [PivotWithColumnReferenceSnippet] + * + * For more information: {@include [DocumentationUrls.Pivot]} + * + * Check out the [`Pivot` Grammar][PivotDocs.Grammar]. + * + * ### Examples + * + * @sample [org.jetbrains.kotlinx.dataframe.samples.api.PivotSamples.withOnPivot_properties] + * @sample [org.jetbrains.kotlinx.dataframe.samples.api.PivotSamples.withOnPivot_strings] + * + * @param [V] The result type of [expression], captured at runtime for constructing the result columns. + * @param [expression] The expression evaluated for each row within each pivot group. + * @return A new [DataRow] with pivot keys as columns or column groups containing the computed results. + */ public inline fun Pivot.with(noinline expression: RowExpression): DataRow = delegate { with(expression) @@ -20,6 +67,25 @@ public inline fun Pivot.with(noinline expression: RowExpressio // region ReducedPivot +/** + * Aggregates this [ReducedPivot] by evaluating the given [row expression][expression] on the single row + * selected by the reducer in each pivot group. + * + * @include [PivotWithReducedResultSnippet] + * + * @include [PivotWithColumnReferenceSnippet] + * + * For more information: {@include [DocumentationUrls.Pivot]} + * + * ### Examples + * + * @sample [org.jetbrains.kotlinx.dataframe.samples.api.PivotSamples.pivotWith_properties] + * @sample [org.jetbrains.kotlinx.dataframe.samples.api.PivotSamples.pivotWith_strings] + * + * @param [V] The result type of [expression], captured at runtime for constructing the result columns. + * @param [expression] The expression evaluated on each selected row. + * @return A new [DataRow] with pivot keys as columns or column groups containing the computed results. + */ public inline fun ReducedPivot.with(noinline expression: RowExpression): DataRow = pivot.delegate { reduce(reducer).with(expression) @@ -29,6 +95,25 @@ public inline fun ReducedPivot.with(noinline expression: RowEx // region PivotGroupBy +/** + * Aggregates this [PivotGroupBy] by evaluating the given [row expression][expression] on every row + * in each group defined by the grouping and pivot keys. + * + * @include [PivotWithOneOrManyResultSnippet] + * + * @include [PivotWithColumnReferenceSnippet] + * + * For more information: {@include [DocumentationUrls.Pivot]} + * + * ### Examples + * + * @sample [org.jetbrains.kotlinx.dataframe.samples.api.PivotSamples.withOnPivotGroupBy] + * + * @param [V] The result type of [expression], captured at runtime for constructing the result columns. + * @param [expression] The expression evaluated for each row within each pivot cell group. + * @return A new [DataFrame] with the grouping key columns, one row per grouping key combination, + * and pivot keys as columns or column groups containing the computed results. + */ public inline fun PivotGroupBy.with(noinline expression: RowExpression): DataFrame { val type = typeOf() return aggregate { internal().withExpr(type, emptyPath(), expression) } @@ -38,6 +123,25 @@ public inline fun PivotGroupBy.with(noinline expression: RowEx // region ReducedPivotGroupBy +/** + * Aggregates this [ReducedPivotGroupBy] by evaluating the given [row expression][expression] on the single row + * selected by the reducer in each group defined by the grouping and pivot keys. + * + * @include [PivotWithReducedResultSnippet] + * + * @include [PivotWithColumnReferenceSnippet] + * + * For more information: {@include [DocumentationUrls.Pivot]} + * + * ### Examples + * + * @sample [org.jetbrains.kotlinx.dataframe.samples.api.PivotSamples.withOnReducedPivotGroupBy] + * + * @param [V] The result type of [expression], captured at runtime for constructing the result columns. + * @param [expression] The expression evaluated on each selected row. + * @return A new [DataFrame] with the grouping key columns, one row per grouping key combination, + * and pivot keys as columns or column groups containing the computed results. + */ public inline fun ReducedPivotGroupBy.with(noinline expression: RowExpression): DataFrame { val type = typeOf() return pivot.aggregate { diff --git a/samples/src/test/kotlin/org/jetbrains/kotlinx/dataframe/samples/api/PivotSamples.kt b/samples/src/test/kotlin/org/jetbrains/kotlinx/dataframe/samples/api/PivotSamples.kt index 0126d5687f..65a1defae3 100644 --- a/samples/src/test/kotlin/org/jetbrains/kotlinx/dataframe/samples/api/PivotSamples.kt +++ b/samples/src/test/kotlin/org/jetbrains/kotlinx/dataframe/samples/api/PivotSamples.kt @@ -514,6 +514,22 @@ class PivotSamples : DataFrameSampleHelper("pivot", "api") { // SampleEnd } + @Test + fun withOnPivotGroupBy() { + // SampleStart + df.groupBy { city }.pivot { isHappy }.with { name.lastName } + df.groupBy("city").pivot("isHappy").with { "name"["lastName"]() } + // SampleEnd + } + + @Test + fun withOnReducedPivotGroupBy() { + // SampleStart + df.groupBy { city }.pivot { isHappy }.maxBy { age }.with { name.lastName } + df.groupBy("city").pivot("isHappy").maxBy("age").with { "name"["lastName"]() } + // SampleEnd + } + // endregion // region aggregate