elasticsearch 9.1.1__py3-none-any.whl → 9.2.0__py3-none-any.whl

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 @@ class ESQL(ABC):
         """
         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.
@@ -143,7 +159,7 @@ class ESQLBase(ABC):
         return False
 
     def change_point(self, value: FieldType) -> "ChangePoint":
-        """`CHANGE_POINT` detects spikes, dips, and change points in a metric.
+        """``CHANGE_POINT`` detects spikes, dips, and change points in a metric.
 
         :param value: The column with the metric in which you want to detect a change point.
 
@@ -163,17 +179,18 @@ class ESQLBase(ABC):
     def completion(
         self, *prompt: ExpressionType, **named_prompt: ExpressionType
     ) -> "Completion":
-        """The `COMPLETION` command allows you to send prompts and context to a Large
+        """The ``COMPLETION`` command allows you to send prompts and context to a Large
         Language Model (LLM) directly within your ES|QL queries, to perform text
         generation tasks.
 
         :param prompt: The input text or expression used to prompt the LLM. This can
                        be a string literal or a reference to a column containing text.
         :param named_prompt: The input text or expresion, given as a keyword argument.
-                             The argument name is used for the column name. If not
-                             specified, the results will be stored in a column named
-                             `completion`. If the specified column already exists, it
-                             will be overwritten with the new results.
+                             The argument name is used for the column name. If the
+                             prompt is given as a positional argument, the results will
+                             be stored in a column named ``completion``. If the
+                             specified column already exists, it will be overwritten
+                             with the new results.
 
         Examples::
 
@@ -283,14 +300,14 @@ class ESQLBase(ABC):
 
     def fork(
         self,
-        fork1: "ESQLBase",
-        fork2: Optional["ESQLBase"] = None,
-        fork3: Optional["ESQLBase"] = None,
-        fork4: Optional["ESQLBase"] = None,
-        fork5: Optional["ESQLBase"] = None,
-        fork6: Optional["ESQLBase"] = None,
-        fork7: Optional["ESQLBase"] = None,
-        fork8: Optional["ESQLBase"] = None,
+        fork1: "Branch",
+        fork2: Optional["Branch"] = None,
+        fork3: Optional["Branch"] = None,
+        fork4: Optional["Branch"] = None,
+        fork5: Optional["Branch"] = None,
+        fork6: Optional["Branch"] = None,
+        fork7: Optional["Branch"] = None,
+        fork8: Optional["Branch"] = None,
     ) -> "Fork":
         """The ``FORK`` processing command creates multiple execution branches to operate on the
         same input data and combines the results in a single output table.
@@ -313,6 +330,51 @@ class ESQLBase(ABC):
             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.
 
@@ -347,6 +409,58 @@ class ESQLBase(ABC):
         """
         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.
@@ -376,7 +490,7 @@ class ESQLBase(ABC):
         return Limit(self, max_number_of_rows)
 
     def lookup_join(self, lookup_index: IndexType) -> "LookupJoin":
-        """`LOOKUP JOIN` enables you to add data from another index, AKA a 'lookup' index,
+        """``LOOKUP JOIN`` enables you to add data from another index, AKA a 'lookup' index,
         to your ES|QL query results, simplifying data enrichment and analysis workflows.
 
         :param lookup_index: The name of the lookup index. This must be a specific index
@@ -410,7 +524,7 @@ class ESQLBase(ABC):
         return LookupJoin(self, lookup_index)
 
     def mv_expand(self, column: FieldType) -> "MvExpand":
-        """The `MV_EXPAND` processing command expands multivalued columns into one row per
+        """The ``MV_EXPAND`` processing command expands multivalued columns into one row per
         value, duplicating other columns.
 
         :param column: The multivalued column to expand.
@@ -439,6 +553,54 @@ class ESQLBase(ABC):
         """
         return Rename(self, **columns)
 
+    def rerank(self, *query: ExpressionType, **named_query: ExpressionType) -> "Rerank":
+        """The ``RERANK`` command uses an inference model to compute a new relevance score
+        for an initial set of documents, directly within your ES|QL queries.
+
+        :param query: The query text used to rerank the documents. This is typically the
+                      same query used in the initial search.
+        :param named_query: The query text used to rerank the documents, given as a
+                            keyword argument. The argument name is used for the column
+                            name. If the query is given as a positional argument, the
+                            results will be stored in a column named ``_score``. If the
+                            specified column already exists, it will be overwritten with
+                            the new results.
+
+        Examples::
+
+            query1 = (
+                ESQL.from_("books").metadata("_score")
+                .where('MATCH(description, "hobbit")')
+                .sort("_score DESC")
+                .limit(100)
+                .rerank("hobbit").on("description").with_(inference_id="test_reranker")
+                .limit(3)
+                .keep("title", "_score")
+            )
+            query2 = (
+                ESQL.from_("books").metadata("_score")
+                .where('MATCH(description, "hobbit") OR MATCH(author, "Tolkien")')
+                .sort("_score DESC")
+                .limit(100)
+                .rerank(rerank_score="hobbit").on("description", "author").with_(inference_id="test_reranker")
+                .sort("rerank_score")
+                .limit(3)
+                .keep("title", "_score", "rerank_score")
+            )
+            query3 = (
+                ESQL.from_("books").metadata("_score")
+                .where('MATCH(description, "hobbit") OR MATCH(author, "Tolkien")')
+                .sort("_score DESC")
+                .limit(100)
+                .rerank(rerank_score="hobbit").on("description", "author").with_(inference_id="test_reranker")
+                .eval(original_score="_score", _score="rerank_score + original_score")
+                .sort("_score")
+                .limit(3)
+                .keep("title", "original_score", "rerank_score", "_score")
+            )
+        """
+        return Rerank(self, *query, **named_query)
+
     def sample(self, probability: float) -> "Sample":
         """The ``SAMPLE`` command samples a fraction of the table rows.
 
@@ -491,7 +653,7 @@ class ESQLBase(ABC):
         :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.
+        Note that only one of ``expressions`` and ``named_expressions`` must be provided.
 
         Examples::
 
@@ -547,7 +709,7 @@ class ESQLBase(ABC):
 
     def where(self, *expressions: ExpressionType) -> "Where":
         """The ``WHERE`` processing command produces a table that contains all the rows
-        from the input table for which the provided condition evaluates to `true`.
+        from the input table for which the provided condition evaluates to ``true``.
 
         :param expressions: A list of boolean expressions, given as positional arguments.
                             These expressions are combined with an ``AND`` logical operator.
@@ -580,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.
         """
@@ -595,7 +759,7 @@ class From(ESQLBase):
 
     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
@@ -643,6 +807,17 @@ class Show(ESQLBase):
         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.
 
@@ -671,21 +846,22 @@ class ChangePoint(ESQLBase):
         self._pvalue_name: Optional[str] = None
 
     def on(self, key: FieldType) -> "ChangePoint":
-        """Continuation of the `CHANGE_POINT` command.
+        """Continuation of the ``CHANGE_POINT`` command.
 
         :param key: The column with the key to order the values by. If not specified,
-                    `@timestamp` is used.
+                    ``@timestamp`` is used.
         """
         self._key = key
         return self
 
     def as_(self, type_name: str, pvalue_name: str) -> "ChangePoint":
-        """Continuation of the `CHANGE_POINT` command.
+        """Continuation of the ``CHANGE_POINT`` command.
 
         :param type_name: The name of the output column with the change point type.
-                          If not specified, `type` is used.
+                          If not specified, ``type`` is used.
         :param pvalue_name: The name of the output column with the p-value that indicates
-                            how extreme the change point is. If not specified, `pvalue` is used.
+                            how extreme the change point is. If not specified, ``pvalue``
+                            is used.
         """
         self._type_name = type_name
         self._pvalue_name = pvalue_name
@@ -722,10 +898,10 @@ class Completion(ESQLBase):
         self._inference_id: Optional[str] = None
 
     def with_(self, inference_id: str) -> "Completion":
-        """Continuation of the `COMPLETION` command.
+        """Continuation of the ``COMPLETION`` command.
 
         :param inference_id: The ID of the inference endpoint to use for the task. The
-                             inference endpoint must be configured with the completion
+                             inference endpoint must be configured with the ``completion``
                              task type.
         """
         self._inference_id = inference_id
@@ -814,7 +990,7 @@ class Enrich(ESQLBase):
         :param match_field: The match field. ``ENRICH`` uses its value to look for records
                             in the enrich index. If not specified, the match will be
                             performed on the column with the same name as the
-                            `match_field` defined in the enrich policy.
+                            ``match_field`` defined in the enrich policy.
         """
         self._match_field = match_field
         return self
@@ -904,14 +1080,14 @@ class Fork(ESQLBase):
     def __init__(
         self,
         parent: ESQLBase,
-        fork1: ESQLBase,
-        fork2: Optional[ESQLBase] = None,
-        fork3: Optional[ESQLBase] = None,
-        fork4: Optional[ESQLBase] = None,
-        fork5: Optional[ESQLBase] = None,
-        fork6: Optional[ESQLBase] = None,
-        fork7: Optional[ESQLBase] = None,
-        fork8: Optional[ESQLBase] = None,
+        fork1: "Branch",
+        fork2: Optional["Branch"] = None,
+        fork3: Optional["Branch"] = None,
+        fork4: Optional["Branch"] = None,
+        fork5: Optional["Branch"] = None,
+        fork6: Optional["Branch"] = None,
+        fork7: Optional["Branch"] = None,
+        fork8: Optional["Branch"] = None,
     ):
         super().__init__(parent)
         self._branches = [fork1, fork2, fork3, fork4, fork5, fork6, fork7, fork8]
@@ -928,6 +1104,39 @@ class Fork(ESQLBase):
         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: Dict[str, Any] = {}
+
+    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.
 
@@ -991,7 +1200,7 @@ class LookupJoin(ESQLBase):
         self._field: Optional[FieldType] = None
 
     def on(self, field: FieldType) -> "LookupJoin":
-        """Continuation of the `LOOKUP_JOIN` command.
+        """Continuation of the ``LOOKUP JOIN`` command.
 
         :param field: The field to join on. This field must exist in both your current query
                       results and in the lookup index. If the field contains multi-valued
@@ -1046,6 +1255,67 @@ class Rename(ESQLBase):
         return f'RENAME {", ".join([f"{self._format_id(old_name)} AS {self._format_id(new_name)}" for old_name, new_name in self._columns.items()])}'
 
 
+class Rerank(ESQLBase):
+    """Implementation of the ``RERANK`` 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, *query: ExpressionType, **named_query: ExpressionType
+    ):
+        if len(query) + len(named_query) > 1:
+            raise ValueError(
+                "this method requires either one positional or one keyword argument only"
+            )
+        super().__init__(parent)
+        self._query = query
+        self._named_query = named_query
+        self._fields: Optional[Tuple[str, ...]] = None
+        self._inference_id: Optional[str] = None
+
+    def on(self, *fields: str) -> "Rerank":
+        """Continuation of the ``RERANK`` command.
+
+        :param fields: One or more fields to use for reranking. These fields should
+                       contain the text that the reranking model will evaluate.
+        """
+        self._fields = fields
+        return self
+
+    def with_(self, inference_id: str) -> "Rerank":
+        """Continuation of the ``RERANK`` command.
+
+        :param inference_id: The ID of the inference endpoint to use for the task. The
+                             inference endpoint must be configured with the ``rerank``
+                             task type.
+        """
+        self._inference_id = inference_id
+        return self
+
+    def _render_internal(self) -> str:
+        if self._fields is None:
+            raise ValueError(
+                "The rerank command requires one or more fields to rerank on"
+            )
+        if self._inference_id is None:
+            raise ValueError("The completion command requires an inference ID")
+        with_ = {"inference_id": self._inference_id}
+        if self._named_query:
+            column = list(self._named_query.keys())[0]
+            query = list(self._named_query.values())[0]
+            query = f"{self._format_id(column)} = {json.dumps(query)}"
+        else:
+            query = json.dumps(self._query[0])
+        return (
+            f"RERANK {query} "
+            f"ON {', '.join([self._format_id(field) for field in self._fields])} "
+            f"WITH {json.dumps(with_)}"
+        )
+
+
 class Sample(ESQLBase):
     """Implementation of the ``SAMPLE`` processing command.
 
@@ -1090,6 +1360,8 @@ class Stats(ESQLBase):
     in a single expression.
     """
 
+    command_name = "STATS"
+
     def __init__(
         self,
         parent: ESQLBase,
@@ -1105,6 +1377,12 @@ class Stats(ESQLBase):
         self._grouping_expressions: Optional[Tuple[ExpressionType, ...]] = None
 
     def by(self, *grouping_expressions: ExpressionType) -> "Stats":
+        """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
+                                     columns, that column will be ignored.
+        """
         self._grouping_expressions = grouping_expressions
         return self
 
@@ -1116,13 +1394,25 @@ class Stats(ESQLBase):
             ]
         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):