Support SAMPLE clause via new QuerySet.sample(sample_fraction, sample_offset)

Author: asantoni

clickhouse_backend/models/query.py

@@ -31,6 +31,11 @@ def prewhere(self, *args, **kwargs):
         clone._query.add_prewhere(Q(*args, **kwargs))
         return clone
 
+    def sample(self, sample_fraction, sample_offset=None):
+        clone = self._chain()
+        clone._query.add_sample(sample_fraction, sample_offset)
+        return clone
+
     def datetimes(self, field_name, kind, order="ASC", tzinfo=None):
         """
         Return a list of datetime objects representing all available

clickhouse_backend/models/sql/compiler.py

@@ -122,6 +122,8 @@ def as_sql(self, with_limits=True, with_col_aliases=False):
         refcounts_before = self.query.alias_refcount.copy()
         try:
             combinator = self.query.combinator
+            sample_fraction = self.query.sample_fraction
+            sample_offset = self.query.sample_offset
             if compat.dj_ge42:
                 extra_select, order_by, group_by = self.pre_sql_setup(
                     with_col_aliases=with_col_aliases or bool(combinator),
@@ -203,6 +205,13 @@ def as_sql(self, with_limits=True, with_col_aliases=False):
                     result += ["FROM", *from_]
                 params.extend(f_params)
 
+                if sample_fraction:
+                    if sample_offset:
+                        sample_sql = "SAMPLE %s OFFSET %s" % (sample_fraction, sample_offset)
+                    else:
+                        sample_sql = "SAMPLE %s" % sample_fraction
+                    result.append(sample_sql)
+
                 if prewhere:
                     result.append("PREWHERE %s" % prewhere)
                     params.extend(p_params)

clickhouse_backend/models/sql/query.py

@@ -19,6 +19,8 @@ def __init__(self, model, where=query.WhereNode, alias_cols=True):
             super().__init__(model, where, alias_cols)
         self.setting_info = {}
         self.prewhere = query.WhereNode()
+        self.sample_fraction = None
+        self.sample_offset = None
 
     def sql_with_params(self):
         """Choose the right db when database router is used."""
@@ -28,6 +30,8 @@ def clone(self):
         obj = super().clone()
         obj.setting_info = self.setting_info.copy()
         obj.prewhere = self.prewhere.clone()
+        obj.sample_fraction = self.sample_fraction
+        obj.sample_offset = self.sample_offset
         return obj
 
     def explain(self, using, format=None, type=None, **settings):
@@ -36,6 +40,10 @@ def explain(self, using, format=None, type=None, **settings):
         compiler = q.get_compiler(using=using)
         return "\n".join(compiler.explain_query())
 
+    def add_sample(self, sample_fraction, sample_offset):
+        self.sample_fraction = sample_fraction
+        self.sample_offset = sample_offset
+
     def add_prewhere(self, q_object):
         """
         refer add_q