Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

[BEAM-13202] Add Coder to CountIfFn.Accum #16856

Merged
merged 3 commits into from
Feb 16, 2022
Merged

[BEAM-13202] Add Coder to CountIfFn.Accum #16856

Show file tree
Hide file tree
Changes from 1 commit
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
add7bbc
[BEAM-13202] Fix typos on tests names for VarianceFnTest
iemejia Feb 15, 2022
1105c34
[BEAM-13202] Add Coder to CountIfFn.Accum
iemejia Feb 15, 2022
cf75357
[BEAM-13202] Reuse Count transform code since CountIf is a specific case
iemejia Feb 15, 2022
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
60 changes: 19 additions & 41 deletions ...ions/sql/src/main/java/org/apache/beam/sdk/extensions/sql/impl/transform/agg/CountIf.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -17,12 +17,11 @@
*/
package org.apache.beam.sdk.extensions.sql.impl.transform.agg;

import com.google.auto.value.AutoValue;
import java.io.Serializable;
import org.apache.beam.sdk.coders.CannotProvideCoderException;
import org.apache.beam.sdk.coders.Coder;
import org.apache.beam.sdk.coders.CoderRegistry;
import org.apache.beam.sdk.coders.SerializableCoder;
import org.apache.beam.sdk.transforms.Combine;
import org.apache.beam.sdk.transforms.Count;

/**
* Returns the count of TRUE values for expression. Returns 0 if there are zero input rows, or if
Expand All @@ -31,62 +30,41 @@
public class CountIf {
private CountIf() {}

public static CountIfFn combineFn() {
return new CountIf.CountIfFn();
public static Combine.CombineFn<Boolean, ?, Long> combineFn() {
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I suppose technically this is a breaking change. But of course everything implementing SQL is not intended for users. Is sql/impl marked @Internal? (this change still LGTM because it is not actually intended as a user API)

Copy link
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Yes, I hesitated but this is internal and this class is only instantiated by SQL when it registers the built-in aggregators so we should be good. sql/impl is not explicitly marked as Internal but I agree that it should.

return new CountIfFn();
}

public static class CountIfFn extends Combine.CombineFn<Boolean, CountIfFn.Accum, Long> {

@AutoValue
public abstract static class Accum implements Serializable {
abstract boolean isExpressionFalse();

abstract long countIfResult();

static Accum empty() {
return of(true, 0L);
}

static Accum of(boolean isExpressionFalse, long countIfResult) {
return new AutoValue_CountIf_CountIfFn_Accum(isExpressionFalse, countIfResult);
}
}
public static class CountIfFn extends Combine.CombineFn<Boolean, long[], Long> {
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Nice cleanup! At this point, would it be worthwhile to let Count.CountFn be public, so we could just inherit it? I don't really have a strong opinion.

Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

IMO using it via composition like this is preferable anyhow. You can see it is more flexible, since you can achieve all the same things without needing it to be public.

private final Combine.CombineFn<Boolean, long[], Long> countFn =
(Combine.CombineFn<Boolean, long[], Long>) Count.<Boolean>combineFn();
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This cast makes me think that Count.combineFn() worked too hard to hide the accumulator type. It should just expose it.

Copy link
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Agree, I hesitated too make it public and make the CountIfFn inherit it and change just to override the addInput method, what let me puzzled was how to make CompineFn composable so I could just 'filter' and then apply CountFn


@Override
public Accum createAccumulator() {
return Accum.empty();
public long[] createAccumulator() {
return countFn.createAccumulator();
}

@Override
public Accum addInput(Accum accum, Boolean input) {
public long[] addInput(long[] accumulator, Boolean input) {
if (Boolean.TRUE.equals(input)) {
return Accum.of(false, accum.countIfResult() + 1);
countFn.addInput(accumulator, input);
}
return accum;
return accumulator;
}

@Override
public Accum mergeAccumulators(Iterable<Accum> accums) {
CountIfFn.Accum merged = createAccumulator();
for (CountIfFn.Accum accum : accums) {
if (!accum.isExpressionFalse()) {
merged = Accum.of(false, merged.countIfResult() + accum.countIfResult());
}
}
return merged;
public long[] mergeAccumulators(Iterable<long[]> accumulators) {
return countFn.mergeAccumulators(accumulators);
}

@Override
public Coder<Accum> getAccumulatorCoder(CoderRegistry registry, Coder<Boolean> inputCoder) {
return SerializableCoder.of(Accum.class);
public Long extractOutput(long[] accumulator) {
return countFn.extractOutput(accumulator);
}

@Override
public Long extractOutput(Accum accum) {
if (!accum.isExpressionFalse()) {
return accum.countIfResult();
}
return 0L;
public Coder<long[]> getAccumulatorCoder(CoderRegistry registry, Coder<Boolean> inputCoder)
throws CannotProvideCoderException {
return countFn.getAccumulatorCoder(registry, inputCoder);
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Technically changing coders here would break in-place update. But SQL really just cannot be relied on for that, since the optimizer might change. So I am just noting that I explicitly say it is OK to break in-place update here.

Copy link
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Technically we did not have a specified Coder so it was breaking when running on a distributed system as the JIRA ticket reported so backwards compatibility seems less of an issue ;)

}
}
}
37 changes: 20 additions & 17 deletions .../sql/src/test/java/org/apache/beam/sdk/extensions/sql/impl/transform/agg/CountIfTest.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -23,53 +23,56 @@
import java.util.Arrays;
import java.util.List;
import org.apache.beam.sdk.coders.BooleanCoder;
import org.apache.beam.sdk.coders.CannotProvideCoderException;
import org.apache.beam.sdk.coders.CoderRegistry;
import org.apache.beam.sdk.transforms.Combine;
import org.junit.Test;

/** Unit tests for {@link CountIf}. */
public class CountIfTest {

@Test
public void testCreatesEmptyAccumulator() {
assertEquals(CountIf.CountIfFn.Accum.empty(), CountIf.combineFn().createAccumulator());
long[] accumulator = (long[]) CountIf.combineFn().createAccumulator();

assertEquals(0, accumulator[0]);
}

@Test
public void testReturnsAccumulatorUnchangedForNullInput() {
CountIf.CountIfFn countIfFn = new CountIf.CountIfFn();
CountIf.CountIfFn.Accum accumulator = countIfFn.createAccumulator();
assertEquals(accumulator, countIfFn.addInput(accumulator, null));
Combine.CombineFn countIfFn = CountIf.combineFn();
long[] accumulator = (long[]) countIfFn.addInput(countIfFn.createAccumulator(), null);

assertEquals(0L, accumulator[0]);
}

@Test
public void testAddsInputToAccumulator() {
CountIf.CountIfFn countIfFn = new CountIf.CountIfFn();
CountIf.CountIfFn.Accum expectedAccumulator = CountIf.CountIfFn.Accum.of(false, 1);
Combine.CombineFn countIfFn = CountIf.combineFn();
long[] accumulator = (long[]) countIfFn.addInput(countIfFn.createAccumulator(), Boolean.TRUE);

assertEquals(
expectedAccumulator, countIfFn.addInput(countIfFn.createAccumulator(), Boolean.TRUE));
assertEquals(1L, accumulator[0]);
}

@Test
public void testCreatesAccumulatorCoder() {
public void testCreatesAccumulatorCoder() throws CannotProvideCoderException {
assertNotNull(
CountIf.combineFn().getAccumulatorCoder(CoderRegistry.createDefault(), BooleanCoder.of()));
}

@Test
public void testMergeAccumulators() {
CountIf.CountIfFn countIfFn = new CountIf.CountIfFn();
List<CountIf.CountIfFn.Accum> accums =
Arrays.asList(CountIf.CountIfFn.Accum.of(false, 2), CountIf.CountIfFn.Accum.of(false, 2));
CountIf.CountIfFn.Accum expectedAccumulator = CountIf.CountIfFn.Accum.of(false, 4);
Combine.CombineFn countIfFn = CountIf.combineFn();
List<long[]> accums = Arrays.asList(new long[] {2}, new long[] {2});
long[] accumulator = (long[]) countIfFn.mergeAccumulators(accums);

assertEquals(expectedAccumulator, countIfFn.mergeAccumulators(accums));
assertEquals(4L, accumulator[0]);
}

@Test
public void testExtractsOutput() {
CountIf.CountIfFn countIfFn = new CountIf.CountIfFn();
CountIf.CountIfFn.Accum expectedAccumulator = countIfFn.createAccumulator();
assertEquals(Long.valueOf(0), countIfFn.extractOutput(expectedAccumulator));
Combine.CombineFn countIfFn = CountIf.combineFn();

assertEquals(0L, countIfFn.extractOutput(countIfFn.createAccumulator()));
}
}