[functional] Add groupby and Stream.groupBy

This commit is contained in:
aNNiMON 2023-10-25 21:36:56 +03:00 committed by Victor Melnik
parent 7a64d7f99d
commit c3a893ea25
5 changed files with 108 additions and 8 deletions

View File

@ -31,6 +31,7 @@ class StreamValue extends MapValue {
set("reduce", wrapTerminal(new functional_reduce()));
set("forEach", wrapTerminal(new functional_forEach()));
set("forEachIndexed", wrapTerminal(new functional_forEachIndexed()));
set("groupBy", wrapTerminal(new functional_groupBy()));
set("toArray", args -> container);
set("joining", container::joinToString);
set("count", args -> NumberValue.of(container.size()));

View File

@ -27,6 +27,7 @@ public final class functional implements Module {
result.put("sortby", new functional_sortBy());
result.put("takewhile", new functional_takeWhile());
result.put("dropwhile", new functional_dropWhile());
result.put("groupby", new functional_groupBy());
result.put("chain", new functional_chain());
result.put("stream", new functional_stream());

View File

@ -0,0 +1,67 @@
package com.annimon.ownlang.modules.functional;
import com.annimon.ownlang.exceptions.TypeException;
import com.annimon.ownlang.lib.*;
import java.util.ArrayList;
import java.util.LinkedHashMap;
import java.util.List;
import java.util.Map;
public final class functional_groupBy implements Function {
@Override
public Value execute(Value[] args) {
Arguments.check(2, args.length);
final Value container = args[0];
final Function classifier = ValueUtils.consumeFunction(args[1], 1);
return groupBy(container, classifier);
}
static Value groupBy(Value container, Function classifier) {
if (container.type() == Types.ARRAY) {
return groupByArray((ArrayValue) container, classifier);
}
if (container.type() == Types.MAP) {
return groupByMap((MapValue) container, classifier);
}
throw new TypeException("Invalid first argument. Array or map expected");
}
@SuppressWarnings("Java8MapApi")
static Value groupByArray(ArrayValue array, Function classifier) {
final var result = new LinkedHashMap<Value, List<Value>>();
for (Value element : array) {
final var key = classifier.execute(element);
var container = result.get(key);
if (container == null) {
container = new ArrayList<>();
result.put(key, container);
}
container.add(element);
}
return fromMapOfArrays(result);
}
static Value groupByMap(MapValue map, Function classifier) {
final var result = new LinkedHashMap<Value, Value>();
for (Map.Entry<Value, Value> element : map) {
final var k = element.getKey();
final var v = element.getValue();
final var key = classifier.execute(k, v);
var container = (MapValue) result.get(key);
if (container == null) {
container = new MapValue(10);
result.put(key, container);
}
container.set(k, v);
}
return new MapValue(result);
}
private static MapValue fromMapOfArrays(Map<Value, List<Value>> map) {
final var result = new LinkedHashMap<Value, Value>();
map.forEach((key, value) -> result.put(key, new ArrayValue(value)));
return new MapValue(result);
}
}

View File

@ -0,0 +1,15 @@
use std, functional
def testArraysGroupBy() {
arr = [1, 2, 3, 4, 1, 2, 3, 1, 2, 3]
result = groupby(arr, def(v) = v % 2 == 0)
assertEquals([2, 4, 2, 2], result[true])
assertEquals([1, 3, 1, 3, 1, 3], result[false])
}
def testMapsGroupBy() {
map = {"abc": 123, "test1": 234, "test2": 345, "test3": 456, "def": 567}
result = groupby(map, def(k, v) = k.startsWith("test"))
assertEquals({"test1": 234, "test2": 345, "test3": 456}, result[true])
assertEquals({"abc": 123, "def": 567}, result[false])
}

View File

@ -53,6 +53,15 @@ def testCustom() {
assertEquals([5,6,4,2], stream(data).custom(::reverse).toArray())
}
def reverse(container) {
size = length(container)
result = newarray(size)
for i : range(size) {
result[size - i - 1] = container[i]
}
return result
}
def testJoining() {
data = [1,2,3,4]
assertEquals("1234", stream(data).joining())
@ -101,11 +110,18 @@ def testForEachMapIndexed() {
assertEquals("a10b21", result)
}
def reverse(container) {
size = length(container)
result = newarray(size)
for i : range(size) {
result[size - i - 1] = container[i]
}
return result
}
def testArraysGroupBy() {
data = [1, 2, 3, 4, 1, 2, 3, 1, 2, 3]
result = stream(data)
.groupBy(def(v) = v % 2 == 0)
assertEquals([2, 4, 2, 2], result[true])
assertEquals([1, 3, 1, 3, 1, 3], result[false])
}
def testMapsGroupBy() {
data = {"abc": 123, "test1": 234, "test2": 345, "test3": 456, "def": 567}
result = stream(data)
.groupBy(def(entry) = entry[0].startsWith("test"))
assertEquals([["test1", 234], ["test2", 345], ["test3", 456]], sort(result[true]))
assertEquals([["abc", 123], ["def", 567]], sort(result[false]))
}