pandas - 将 stream_results=True (execution_options) 添加到 kedro.extras.datasets.pandas.SQLQueryDataSet

Question

是否可以将 execution_options 添加到 kedro.extras.datasets.pandas.SQLQueryDataSet？

例如，我想将 stream_results=True 添加到连接字符串中。

engine = create_engine("postgresql://postgres:pass@localhost/example") conn = engine.connect().execution_options(stream_results=True)

这是我的目录.yml

table_name:
  type: pandas.SQLQueryDataSet
  credentials: creds
  sql: select * from table
  load_args:
    chunksize: 1000

关于如何使用 pandas.SQLQueryDataSet 添加/编辑 execution_options 的任何想法？具体来说，stream_results=True。

Answer 1

您可能需要在现有的之上创建一个薄层SQLQueryDataSet：

class CustomSQLQueryDataSet(kedro.extras.datasets.SQLQueryDataSet):
  def _load(self, *args, **kwargs):
      self._load_args["con"] = create_engine(self._load_args["con"]).connect().execution_options(stream_results=True)
      return super()._load(*args, **kwargs)

然后在你的目录中使用这个类。