轻松将 ES|QL 查询结果转换为 Python Pandas dataframe

原创

点火三周

发布于 2024-04-14 08:16:15

2090

发布于 2024-04-14 08:16:15

Elasticsearch 查询语言（ES|QL）为我们提供了一种强大的方式，用于过滤、转换和分析存储在 Elasticsearch 中的数据。它设计简单易学易用，非常适合熟悉 Pandas 和其他基于数据框的库的数据科学家。实际上，ES|QL 查询产生的表格具有命名列，这就是数据框的定义！

ES|QL 生成表格

首先，让我们导入一些测试数据。我们将使用员工样本数据和映射。加载这个数据集的最简单方法是在 Kibana 控制台中运行这两个 Elasticsearch API 请求。

好的，既然这个环节已经完成，让我们使用 ES|QL CSV 导出功能，将完整的员工数据集转换为 Pandas DataFrame 对象：

from io import StringIO
from elasticsearch import Elasticsearch
import pandas as pd

client = Elasticsearch(
    "https://[host].elastic-cloud.com",
    api_key="...",
)

response = client.esql.query(
    query="FROM employees | LIMIT 500",
    format="csv",
)
df = pd.read_csv(StringIO(response.body))
print(df)

尽管这个数据集只包含 100 条记录，但我们使用 LIMIT 命令是为了避免 ES|QL 警告我们可能丢失记录。这将打印出如下数据框：

    avg_worked_seconds  ...  salary_change.long still_hired
0            268728049  ...                   1        True
1            328922887  ...            [-7, 11]        True
2            200296405  ...            [12, 14]       False
3            311267831  ...       [0, 1, 3, 13]        True
4            244294991  ...            [-2, 13]        True
..                 ...  ...                 ...         ...
95           204381503  ...                 NaN       False
96           206258084  ...                  -1       False
97           272392146  ...          [-2, 4, 8]       False
98           377713748  ...    [-8, -3, 10, 14]        True
99           223910853  ...            [-7, 13]        True

这意味着您现在可以使用 Pandas 分析数据。但您也可以继续使用 ES|QL 处理数据，这在查询返回超过 10,000 行时特别有用，这是 ES|QL 查询可以返回的最大行数。

在下一个示例中，我们通过使用 STATS ... BY（类似于 SQL 中的 GROUP BY）来统计说某种语言的员工数量。然后我们使用 SORT 对结果进行语言列排序：

response = client.esql.query(
    query="""
    FROM employees
    | STATS count = COUNT(emp_no) BY languages
    | SORT languages
    | LIMIT 500
    """,
    format="csv",
)

df = pd.read_csv(
    StringIO(response.body),
    dtype={"count": "Int64", "languages": "Int64"},
)
print(df)

请注意，我们在这里使用了 pd.read_csv() 的 dtype 参数，这在 Pandas 推断的类型不够时非常有用。上述代码打印出以下结果：

   count  languages
0     15          1
1     19          2
2     17          3
3     18          4
4     21          5

21 名员工会说 5 种语言，哇！

最后，假设您的代码的最终用户可以控制说话的最低语言数量。您可以直接在 Python 中格式化查询，但这将允许攻击者执行 ES|QL 注入！相反，使用 ES|QL REST API 内置的参数支持功能：

response = client.esql.query(
    query="""
    FROM employees
    | STATS count = COUNT(emp_no) BY languages
    | WHERE languages >= (?)
    | SORT languages
    | LIMIT 500
    """,
    format="csv",
    params=[3],
)

df = pd.read_csv(
    StringIO(response.body),
    dtype={"count": "Int64", "languages": "Int64"},
)
print(df)

这将打印出以下结果：

   count  languages
0     17          3
1     18          4
2     21          5

如您所见，ES|QL 和 Pandas 可以很好地协同工作。然而，CSV 并不是理想的格式，因为它需要显式类型声明，并且对 ES|QL 产生的一些更复杂的结果（如嵌套数组和对象）处理不佳。为此，我们正在努力为 ES|QL 添加对 Apache Arrow 数据框的原生支持，这将使所有这些变得透明，并带来显著的性能提升。

如果您想了解更多关于 ES|QL 的信息，ES|QL 文档是最佳起点。要了解更多关于 Python Elasticsearch 客户端的信息，您可以查阅文档，在 Discuss 上用 language-clients 标签提问，或者如果您发现了一个错误或有功能请求，可以打开一个新问题。谢谢！

原创声明：本文系作者授权腾讯云开发者社区发表，未经许可，不得转载。

如有侵权，请联系 cloudcommunity@tencent.com 删除。

Elasticsearch Service