add TS, FUSE and INLINE STATS commands

Author: miguelgrinberg

elasticsearch/esql/esql.py

@@ -18,7 +18,7 @@
 import json
 import re
 from abc import ABC, abstractmethod
-from typing import Any, Dict, Optional, Tuple, Type, Union
+from typing import Any, Dict, List, Optional, Tuple, Type, Union
 
 from ..dsl.document_base import DocumentBase, InstrumentedExpression, InstrumentedField
 
@@ -78,6 +78,22 @@ def show(item: str) -> "Show":
         """
         return Show(item)
 
+    @staticmethod
+    def ts(*indices: IndexType) -> "TS":
+        """The ``TS`` source command is similar to ``FROM``, but for time series indices.
+
+        :param indices: A list of indices, data streams or aliases. Supports wildcards and date math.
+
+        Examples::
+
+            query = (
+                ESQL.ts("metrics")
+                .where("@timestamp >= now() - 1 day")
+                .stats("SUM(AVG_OVER_TIME(memory_usage)").by("host", "TBUCKET(1 hour)")
+            )
+        """
+        return TS(*indices)
+
     @staticmethod
     def branch() -> "Branch":
         """This method can only be used inside a ``FORK`` command to create each branch.
@@ -314,6 +330,51 @@ def fork(
             raise ValueError("a query can only have one fork")
         return Fork(self, fork1, fork2, fork3, fork4, fork5, fork6, fork7, fork8)
 
+    def fuse(self, method: Optional[str] = None) -> "Fuse":
+        """The ``FUSE`` processing command merges rows from multiple result sets and assigns
+        new relevance scores.
+
+        :param method: Defaults to ``RRF``. Can be one of ``RRF`` (for Reciprocal Rank Fusion)
+                       or ``LINEAR`` (for linear combination of scores). Designates which
+                       method to use to assign new relevance scores.
+
+        Examples::
+
+            query1 = (
+                ESQL.from_("books").metadata("_id", "_index", "_score")
+                .fork(
+                    ESQL.branch().where('title:"Shakespeare"').sort("_score DESC"),
+                    ESQL.branch().where('semantic_title:"Shakespeare"').sort("_score DESC"),
+                )
+                .fuse()
+            )
+            query2 = (
+                ESQL.from_("books").metadata("_id", "_index", "_score")
+                .fork(
+                    ESQL.branch().where('title:"Shakespeare"').sort("_score DESC"),
+                    ESQL.branch().where('semantic_title:"Shakespeare"').sort("_score DESC"),
+                )
+                .fuse("linear")
+            )
+            query3 = (
+                ESQL.from_("books").metadata("_id", "_index", "_score")
+                .fork(
+                    ESQL.branch().where('title:"Shakespeare"').sort("_score DESC"),
+                    ESQL.branch().where('semantic_title:"Shakespeare"').sort("_score DESC"),
+                )
+                .fuse("linear").by("title", "description")
+            )
+            query4 = (
+                ESQL.from_("books").metadata("_id", "_index", "_score")
+                .fork(
+                    ESQL.branch().where('title:"Shakespeare"').sort("_score DESC"),
+                    ESQL.branch().where('semantic_title:"Shakespeare"').sort("_score DESC"),
+                )
+                .fuse("linear").with_(normalizer="minmax")
+            )
+        """
+        return Fuse(self, method)
+
     def grok(self, input: FieldType, pattern: str) -> "Grok":
         """``GROK`` enables you to extract structured data out of a string.
 
@@ -348,6 +409,58 @@ def grok(self, input: FieldType, pattern: str) -> "Grok":
         """
         return Grok(self, input, pattern)
 
+    def inline_stats(
+        self, *expressions: ExpressionType, **named_expressions: ExpressionType
+    ) -> "Stats":
+        """The ``INLINE STATS`` processing command groups rows according to a common value
+        and calculates one or more aggregated values over the grouped rows.
+
+        The command is identical to ``STATS`` except that it preserves all the columns from
+        the input table.
+
+        :param expressions: A list of expressions, given as positional arguments.
+        :param named_expressions: A list of expressions, given as keyword arguments. The
+                                  argument names are used for the returned aggregated values.
+
+        Note that only one of ``expressions`` and ``named_expressions`` must be provided.
+
+        Examples::
+
+            query1 = (
+                ESQL.from_("employees")
+                .keep("emp_no", "languages", "salary")
+                .inline_stats(max_salary=functions.max(E("salary"))).by("languages")
+            )
+            query2 = (
+                ESQL.from_("employees")
+                .keep("emp_no", "languages", "salary")
+                .inline_stats(max_salary=functions.max(E("salary")))
+            )
+            query3 = (
+                ESQL.from_("employees")
+                .where("still_hired")
+                .keep("emp_no", "languages", "salary", "hire_date")
+                .eval(tenure=functions.date_diff("year", E("hire_date"), "2025-09-18T00:00:00"))
+                .drop("hire_date")
+                .inline_stats(
+                    avg_salary=functions.avg(E("salary")),
+                    count=functions.count(E("*")),
+                )
+                .by("languages", "tenure")
+            )
+            query4 = (
+                ESQL.from_("employees")
+                .keep("emp_no", "salary")
+                .inline_stats(
+                    avg_lt_50=functions.round(functions.avg(E("salary"))).where(E("salary") < 50000),
+                    avg_lt_60=functions.round(functions.avg(E("salary"))).where(E("salary") >= 50000, E("salary") < 60000),
+                    avg_gt_60=functions.round(functions.avg(E("salary"))).where(E("salary") >= 60000),
+                )
+            )
+
+        """
+        return InlineStats(self, *expressions, **named_expressions)
+
     def keep(self, *columns: FieldType) -> "Keep":
         """The ``KEEP`` processing command enables you to specify what columns are returned
         and the order in which they are returned.
@@ -629,13 +742,15 @@ class From(ESQLBase):
     in a single expression.
     """
 
+    command_name = "FROM"
+
     def __init__(self, *indices: IndexType):
         super().__init__()
         self._indices = indices
         self._metadata_fields: Tuple[FieldType, ...] = tuple()
 
     def metadata(self, *fields: FieldType) -> "From":
-        """Continuation of the ``FROM`` source command.
+        """Continuation of the ``FROM`` and ``TS`` source commands.
 
         :param fields: metadata fields to retrieve, given as positional arguments.
         """
@@ -644,7 +759,7 @@ def metadata(self, *fields: FieldType) -> "From":
 
     def _render_internal(self) -> str:
         indices = [self._format_index(index) for index in self._indices]
-        s = f'{self.__class__.__name__.upper()} {", ".join(indices)}'
+        s = f'{self.command_name} {", ".join(indices)}'
         if self._metadata_fields:
             s = (
                 s
@@ -692,6 +807,17 @@ def _render_internal(self) -> str:
         return f"SHOW {self._format_id(self._item)}"
 
 
+class TS(From):
+    """Implementation of the ``TS`` source command.
+
+    This class inherits from :class:`ESQLBase <elasticsearch.esql.esql.ESQLBase>`,
+    to make it possible to chain all the commands that belong to an ES|QL query
+    in a single expression.
+    """
+
+    command_name = "TS"
+
+
 class Branch(ESQLBase):
     """Implementation of a branch inside a ``FORK`` processing command.
 
@@ -978,6 +1104,39 @@ def _render_internal(self) -> str:
         return f"FORK {cmds}"
 
 
+class Fuse(ESQLBase):
+    """Implementation of the ``FUSE`` processing command.
+
+    This class inherits from :class:`ESQLBase <elasticsearch.esql.esql.ESQLBase>`,
+    to make it possible to chain all the commands that belong to an ES|QL query
+    in a single expression.
+    """
+
+    def __init__(self, parent: ESQLBase, method: Optional[str] = None):
+        super().__init__(parent)
+        self.method = method
+        self.by_columns: List[FieldType] = []
+        self.options = {}
+
+    def by(self, *columns: FieldType) -> "Fuse":
+        self.by_columns += list(columns)
+        return self
+
+    def with_(self, **options: Any) -> "Fuse":
+        self.options = options
+        return self
+
+    def _render_internal(self) -> str:
+        method = f" {self.method.upper()}" if self.method else ""
+        by = (
+            " " + " ".join([f"BY {column}" for column in self.by_columns])
+            if self.by_columns
+            else ""
+        )
+        with_ = " WITH " + json.dumps(self.options) if self.options else ""
+        return f"FUSE{method}{by}{with_}"
+
+
 class Grok(ESQLBase):
     """Implementation of the ``GROK`` processing command.
 
@@ -1201,6 +1360,8 @@ class Stats(ESQLBase):
     in a single expression.
     """
 
+    command_name = "STATS"
+
     def __init__(
         self,
         parent: ESQLBase,
@@ -1216,7 +1377,7 @@ def __init__(
         self._grouping_expressions: Optional[Tuple[ExpressionType, ...]] = None
 
     def by(self, *grouping_expressions: ExpressionType) -> "Stats":
-        """Continuation of the ``STATS`` command.
+        """Continuation of the ``STATS`` and ``INLINE STATS`` commands.
 
         :param grouping_expressions: Expressions that output the values to group by.
                                      If their names coincide with one of the computed
@@ -1233,13 +1394,25 @@ def _render_internal(self) -> str:
             ]
         else:
             exprs = [f"{self._format_expr(expr)}" for expr in self._expressions]
-        expression_separator = ",\n        "
+        indent = " " * (len(self.command_name) + 3)
+        expression_separator = f",\n{indent}"
         by = (
             ""
             if self._grouping_expressions is None
-            else f'\n        BY {", ".join([f"{self._format_expr(expr)}" for expr in self._grouping_expressions])}'
+            else f'\n{indent}BY {", ".join([f"{self._format_expr(expr)}" for expr in self._grouping_expressions])}'
         )
-        return f'STATS {expression_separator.join([f"{expr}" for expr in exprs])}{by}'
+        return f'{self.command_name} {expression_separator.join([f"{expr}" for expr in exprs])}{by}'
+
+
+class InlineStats(Stats):
+    """Implementation of the ``INLINE STATS`` processing command.
+
+    This class inherits from :class:`ESQLBase <elasticsearch.esql.esql.ESQLBase>`,
+    to make it possible to chain all the commands that belong to an ES|QL query
+    in a single expression.
+    """
+
+    command_name = "INLINE STATS"
 
 
 class Where(ESQLBase):