<commit_before>import sqlite3
import string
import time
from operator import attrgetter
from typing import Callable, Dict, List, NamedTuple, Optional, Tuple

from numba import float64, int64, optional
from numba.experimental import jitclass
from numba.types import ClassType

from slumba import create_aggregate, sqlite_udaf


@sqlite_udaf(float64(float64))
@jitclass(
    [
        ("mean", float64),
        ("sum_of_squares_of_differences", float64),
        ("count", int64),
    ]
)
class Var:
    def __init__(self):
        self.mean = 0.0
        self.sum_of_squares_of_differences = 0.0
        self.count = 0

    def step(self, value: float) -> None:
        self.count += 1
        delta = value - self.mean
        self.mean += delta
        self.sum_of_squares_of_differences += delta * (value - self.mean)

    def finalize(self) -> float:
        return self.sum_of_squares_of_differences / (self.count - 1)


@sqlite_udaf(optional(float64)(optional(float64), optional(float64)))
@jitclass(
    [("mean1", float64), ("mean2", float64), ("mean12", float64), ("count", int64)]
)
class Cov:
    def __init__(self):
        self.mean1 = 0.0
        self.mean2 = 0.0
        self.mean12 = 0.0
        self.count = 0

    def step(self, x: Optional[float], y: Optional[float]) -> None:
        if x is not None and y is not None:
            self.count += 1
            n = self.count
            delta1 = (x - self.mean1) / n
            self.mean1 += delta1
            delta2 = (y - self.mean2) / n
            self.mean2 += delta2
            self.mean12 += (n - 1) * delta1 * delta2 - self.mean12 / n

    def finalize(self) -> Optional[float]:
        n = self.count
        if not n:
            return None
        return n / (n - 1) * self.mean12


@sqlite_udaf(optional(float64)(optional(float64)))
@jitclass([("total", float64), ("count", int64)])
class Avg:
    def __init__(self):
        self.total = 0.0
        self.count = 0

    def step(self, value: Optional[float]) -> None:
        if value is not None:
            self.total += value
            self.count += 1

    def finalize(self) -> Optional[float]:
        if not self.count:
            return None
        return self.total / self.count


@sqlite_udaf(optional(float64)(float64))
@jitclass([("total", float64), ("count", int64)])
class Sum:
    def __init__(self):
        self.total = 0.0
        self.count = 0

    def step(self, value: Optional[float]) -> None:
        if value is not None:
            self.total += value
            self.count += 1

    def finalize(self) -> Optional[float]:
        return self.total if self.count > 0 else None<commit_msg>style: add return type annotations to aggregate.py<commit_after>import sqlite3
import string
import time
from operator import attrgetter
from typing import Callable, Dict, List, NamedTuple, Optional, Tuple

from numba import float64, int64, optional
from numba.experimental import jitclass
from numba.types import ClassType

from slumba import create_aggregate, sqlite_udaf


@sqlite_udaf(float64(float64))
@jitclass(
    [
        ("mean", float64),
        ("sum_of_squares_of_differences", float64),
        ("count", int64),
    ]
)
class Var:
    def __init__(self) -> None:
        self.mean = 0.0
        self.sum_of_squares_of_differences = 0.0
        self.count = 0

    def step(self, value: float) -> None:
        self.count += 1
        delta = value - self.mean
        self.mean += delta
        self.sum_of_squares_of_differences += delta * (value - self.mean)

    def finalize(self) -> float:
        return self.sum_of_squares_of_differences / (self.count - 1)


@sqlite_udaf(optional(float64)(optional(float64), optional(float64)))
@jitclass(
    [("mean1", float64), ("mean2", float64), ("mean12", float64), ("count", int64)]
)
class Cov:
    def __init__(self) -> None:
        self.mean1 = 0.0
        self.mean2 = 0.0
        self.mean12 = 0.0
        self.count = 0

    def step(self, x: Optional[float], y: Optional[float]) -> None:
        if x is not None and y is not None:
            self.count += 1
            n = self.count
            delta1 = (x - self.mean1) / n
            self.mean1 += delta1
            delta2 = (y - self.mean2) / n
            self.mean2 += delta2
            self.mean12 += (n - 1) * delta1 * delta2 - self.mean12 / n

    def finalize(self) -> Optional[float]:
        n = self.count
        if not n:
            return None
        return n / (n - 1) * self.mean12


@sqlite_udaf(optional(float64)(optional(float64)))
@jitclass([("total", float64), ("count", int64)])
class Avg:
    def __init__(self) -> None:
        self.total = 0.0
        self.count = 0

    def step(self, value: Optional[float]) -> None:
        if value is not None:
            self.total += value
            self.count += 1

    def finalize(self) -> Optional[float]:
        if not self.count:
            return None
        return self.total / self.count


@sqlite_udaf(optional(float64)(float64))
@jitclass([("total", float64), ("count", int64)])
class Sum:
    def __init__(self) -> None:
        self.total = 0.0
        self.count = 0

    def step(self, value: Optional[float]) -> None:
        if value is not None:
            self.total += value
            self.count += 1

    def finalize(self) -> Optional[float]:
        return self.total if self.count > 0 else None
{'fraction_non_alphanumeric': 0.1233512911016162, 'fraction_numerical': 0.035296303176667286, 'mean_word_length': 2.263030303030303, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 25, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '4586662', 'n_tokens_mistral': 1855, 'n_tokens_neox': 1713, 'n_words': 473}