8432690087

Committed 26 Mar 2024 08:07AM UTC coverage: 85.986% (+0.04%) from 85.943%

Build # 8432690087

Build Type

push

github

Committed by

web-flow

Commit Message

Merge pull request #303 from aronbierbaum/add_quantile

Add support for quantile and quantileIf functions

Run Details

33 of 37 new or added lines in 2 files covered. (89.19%)

2399 of 2790 relevant lines covered (85.99%)

20.04 hits per line

Source File
Press 'n' to go to next uncovered line, 'b' for previous

88.89

/clickhouse_sqlalchemy/sql/functions.py

from __future__ import annotations

from typing import TYPE_CHECKING, Any, TypeVar

from sqlalchemy.ext.compiler import compiles
from sqlalchemy.sql import coercions, roles
from sqlalchemy.sql.elements import ColumnElement
from sqlalchemy.sql.functions import GenericFunction

from clickhouse_sqlalchemy import types

if TYPE_CHECKING:
    from sqlalchemy.sql._typing import _ColumnExpressionArgument

_T = TypeVar('_T', bound=Any)


class quantile(GenericFunction[_T]):
    inherit_cache = True

    def __init__(
        self, level: float, expr: _ColumnExpressionArgument[Any],
        condition: _ColumnExpressionArgument[Any] = None, **kwargs: Any
    ):
        arg: ColumnElement[Any] = coercions.expect(
            roles.ExpressionElementRole, expr, apply_propagate_attrs=self
        )

        args = [arg]
        if condition is not None:
            condition = coercions.expect(
                roles.ExpressionElementRole, condition,
                apply_propagate_attrs=self
            )
            args.append(condition)

        self.level = level

        if isinstance(arg.type, (types.Decimal, types.Float, types.Int)):
            return_type = types.Float64
        elif isinstance(arg.type, types.DateTime):
            return_type = types.DateTime
        elif isinstance(arg.type, types.Date):
            return_type = types.Date
        else:
            return_type = types.Float64

        kwargs['type_'] = return_type
        kwargs['_parsed_args'] = args
        super().__init__(arg, **kwargs)


class quantileIf(quantile[_T]):
    inherit_cache = True


@compiles(quantile, 'clickhouse')
@compiles(quantileIf, 'clickhouse')
def compile_quantile(element, compiler, **kwargs):
    args_str = compiler.function_argspec(element, **kwargs)
    return f'{element.name}({element.level}){args_str}'

1	from __future__ import annotations	16✔
2
3	from typing import TYPE_CHECKING, Any, TypeVar	24✔
4
5	from sqlalchemy.ext.compiler import compiles	24✔
6	from sqlalchemy.sql import coercions, roles	24✔
7	from sqlalchemy.sql.elements import ColumnElement	24✔
8	from sqlalchemy.sql.functions import GenericFunction	24✔
9
10	from clickhouse_sqlalchemy import types	24✔
11
12	if TYPE_CHECKING:	24✔
NEW 13	from sqlalchemy.sql._typing import _ColumnExpressionArgument	×
14
15	_T = TypeVar('_T', bound=Any)	24✔
16
17
18	class quantile(GenericFunction[_T]):	24✔
19	inherit_cache = True	24✔
20
21	def __init__(	24✔
22	self, level: float, expr: _ColumnExpressionArgument[Any],
23	condition: _ColumnExpressionArgument[Any] = None, **kwargs: Any
24	):
25	arg: ColumnElement[Any] = coercions.expect(	24✔
26	roles.ExpressionElementRole, expr, apply_propagate_attrs=self
27	)
28
29	args = [arg]	24✔
30	if condition is not None:	24✔
31	condition = coercions.expect(	24✔
32	roles.ExpressionElementRole, condition,
33	apply_propagate_attrs=self
34	)
35	args.append(condition)	24✔
36
37	self.level = level	24✔
38
39	if isinstance(arg.type, (types.Decimal, types.Float, types.Int)):	24✔
40	return_type = types.Float64	24✔
41	elif isinstance(arg.type, types.DateTime):	24✔
42	return_type = types.DateTime	24✔
NEW 43	elif isinstance(arg.type, types.Date):	×
NEW 44	return_type = types.Date	×
45	else:
NEW 46	return_type = types.Float64	×
47
48	kwargs['type_'] = return_type	24✔
49	kwargs['_parsed_args'] = args	24✔
50	super().__init__(arg, **kwargs)	24✔
51
52
53	class quantileIf(quantile[_T]):	24✔
54	inherit_cache = True	24✔
55
56
57	@compiles(quantile, 'clickhouse')	24✔
58	@compiles(quantileIf, 'clickhouse')	24✔
59	def compile_quantile(element, compiler, **kwargs):	20✔
60	args_str = compiler.function_argspec(element, **kwargs)	24✔
61	return f'{element.name}({element.level}){args_str}'	24✔

xzkostyan / clickhouse-sqlalchemy / 8432690087

Source File Press 'n' to go to next uncovered line, 'b' for previous

Source File
Press 'n' to go to next uncovered line, 'b' for previous