#@title Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# https://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""
Explanation: Copyright 2020 The TensorFlow IO Authors.
End of explanation
"""


try:
  %tensorflow_version 2.x
except Exception:
  pass

!pip install tensorflow-io

"""
Explanation: 从 TensorFlow IO 中读取 PostgreSQL 数据库
<table class="tfo-notebook-buttons" align="left">
  <td><a target="_blank" href="https://tensorflow.google.cn/io/tutorials/postgresql"><img src="https://tensorflow.google.cn/images/tf_logo_32px.png">在 TensorFlow.org 上查看 </a></td>
  <td><a target="_blank" href="https://colab.research.google.com/github/tensorflow/docs-l10n/blob/master/site/zh-cn/io/tutorials/postgresql.ipynb"><img src="https://tensorflow.google.cn/images/colab_logo_32px.png">在 Google Colab 中运行 </a></td>
  <td><a target="_blank" href="https://github.com/tensorflow/docs-l10n/blob/master/site/zh-cn/io/tutorials/postgresql.ipynb"><img src="https://tensorflow.google.cn/images/GitHub-Mark-32px.png">在 GitHub 中查看源代码</a></td>
      <td><a href="https://storage.googleapis.com/tensorflow_docs/docs-l10n/site/zh-cn/io/tutorials/postgresql.ipynb"><img src="https://tensorflow.google.cn/images/download_logo_32px.png">下载笔记本</a></td>
</table>

概述
本教程介绍了如何从 PostgreSQL 数据库服务器创建 tf.data.Dataset，以便可以将创建的 Dataset 传递给 tf.keras 进行训练或推理。
SQL 数据库是数据科学家的重要数据来源。PostgreSQL 是最流行的开源 SQL 数据库之一，在企业中广泛用于全面存储关键数据和事务数据。直接从 PostgreSQL 数据库服务器创建 Dataset 并将 Dataset 传递给 tf.keras 进行训练或推理，可以极大地简化数据流水线，并帮助数据科学家专注于构建机器学习模型。
设置和用法
安装所需的 tensorflow-io 软件包，然后重新启动运行时
End of explanation
"""


# Install postgresql server
!sudo apt-get -y -qq update
!sudo apt-get -y -qq install postgresql
!sudo service postgresql start

# Setup a password `postgres` for username `postgres`
!sudo -u postgres psql -U postgres -c "ALTER USER postgres PASSWORD 'postgres';"

# Setup a database with name `tfio_demo` to be used
!sudo -u postgres psql -U postgres -c 'DROP DATABASE IF EXISTS tfio_demo;'
!sudo -u postgres psql -U postgres -c 'CREATE DATABASE tfio_demo;'

"""
Explanation: 安装和设置 PostgreSQL（可选）
警告：此笔记本只适合在 Google Colab 中运行。它会在系统上安装软件包，并且需要 sudo 访问。如果要在本地 Jupyter 笔记本中运行，请谨慎操作。
为了演示在 Google Colab 上的用法，您将安装 PostgreSQL 服务器。此外，还需要密码和一个空数据库。
如果您不在 Google Colab 上运行此笔记本，或者希望使用现有数据库，请跳过以下设置并前进到下一部分。
End of explanation
"""


%env TFIO_DEMO_DATABASE_NAME=tfio_demo
%env TFIO_DEMO_DATABASE_HOST=localhost
%env TFIO_DEMO_DATABASE_PORT=5432
%env TFIO_DEMO_DATABASE_USER=postgres
%env TFIO_DEMO_DATABASE_PASS=postgres

"""
Explanation: 设置必要的环境变量
以下环境变量基于上一部分中的 PostgreSQL 设置。如果您使用其他设置或使用现有数据库，则应相应地更改这些设置：
End of explanation
"""


!curl -s -OL https://github.com/tensorflow/io/raw/master/docs/tutorials/postgresql/AirQualityUCI.sql

!PGPASSWORD=$TFIO_DEMO_DATABASE_PASS psql -q -h $TFIO_DEMO_DATABASE_HOST -p $TFIO_DEMO_DATABASE_PORT -U $TFIO_DEMO_DATABASE_USER -d $TFIO_DEMO_DATABASE_NAME -f AirQualityUCI.sql

"""
Explanation: 在 PostgreSQL 服务器中准备数据
出于演示目的，本教程将创建一个数据库，并使用一些数据填充该数据库。本教程中使用的数据来自 UCI Machine Learning Repository 中提供的 Air Quality Data Set。
下面是 Air Quality Data Set 的一个子集的预览：
日期 | 时间 | CO(GT) | PT08.S1(CO) | NMHC(GT) | C6H6(GT) | PT08.S2(NMHC) | NOx(GT) | PT08.S3(NOx) | NO2(GT) | PT08.S4(NO2) | PT08.S5(O3) | T | RH | AH
--- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | ---
10/03/2004 | 18.00.00 | 2.6 | 1360 | 150 | 11.9 | 1046 | 166 | 1056 | 113 | 1692 | 1268 | 13.6 | 48.9 | 0.7578
10/03/2004 | 19.00.00 | 2 | 1292 | 112 | 9.4 | 955 | 103 | 1174 | 92 | 1559 | 972 | 13.3 | 47.7 | 0.7255
10/03/2004 | 20.00.00 | 2.2 | 1402 | 88 | 9.0 | 939 | 131 | 1140 | 114 | 1555 | 1074 | 11.9 | 54.0 | 0.7502
10/03/2004 | 21.00.00 | 2.2 | 1376 | 80 | 9.2 | 948 | 172 | 1092 | 122 | 1584 | 1203 | 11.0 | 60.0 | 0.7867
10/03/2004 | 22.00.00 | 1.6 | 1272 | 51 | 6.5 | 836 | 131 | 1205 | 116 | 1490 | 1110 | 11.2 | 59.6 | 0.7888
参考文献部分提供了有关 Air Quality Data Set 和 UCI Machine Learning Repository 的更多信息。
为了帮助简化数据准备工作，我们准备了 sql 版本的 Air Quality Data Set，并将其以 AirQualityUCI.sql 形式提供。
创建表的语句如下：
CREATE TABLE AirQualityUCI (
  Date DATE,
  Time TIME,
  CO REAL,
  PT08S1 INT,
  NMHC REAL,
  C6H6 REAL,
  PT08S2 INT,
  NOx REAL,
  PT08S3 INT,
  NO2 REAL,
  PT08S4 INT,
  PT08S5 INT,
  T REAL,
  RH REAL,
  AH REAL
);
用于在数据库中创建表并填充数据的完整命令如下：
End of explanation
"""


import os
import tensorflow_io as tfio

endpoint="postgresql://{}:{}@{}?port={}&dbname={}".format(
    os.environ['TFIO_DEMO_DATABASE_USER'],
    os.environ['TFIO_DEMO_DATABASE_PASS'],
    os.environ['TFIO_DEMO_DATABASE_HOST'],
    os.environ['TFIO_DEMO_DATABASE_PORT'],
    os.environ['TFIO_DEMO_DATABASE_NAME'],
)

dataset = tfio.experimental.IODataset.from_sql(
    query="SELECT co, pt08s1 FROM AirQualityUCI;",
    endpoint=endpoint)

print(dataset.element_spec)

"""
Explanation: 从 PostgreSQL 服务器创建数据集并在 TensorFlow 中使用
从 PostgreSQL 服务器创建数据集就像调用带 query 和 endpoint 参数的 tfio.experimental.IODataset.from_sql 一样简单。query 是对表中选定列的 SQL 查询，而 endpoint 参数是地址和数据库名称：
End of explanation
"""


dataset = tfio.experimental.IODataset.from_sql(
    query="SELECT nox, no2 FROM AirQualityUCI;",
    endpoint=endpoint)

dataset = dataset.map(lambda e: (e['nox'] - e['no2']))

# check only the first 20 record
dataset = dataset.take(20)

print("NOx - NO2:")
for difference in dataset:
  print(difference.numpy())

"""
Explanation: 从上面的 dataset.element_spec 输出可以看出，所创建的 Dataset 的元素是一个 python dict 对象，其中以数据库表的列名作为键。这对于应用进一步运算非常方便。例如，可以同时选择 Dataset 的 nox 和 no2 字段，并计算差值：
End of explanation
"""


