DuckDB plugin

[samhita-alla]

Signed-off-by: Samhita Alla aallasamhita@gmail.com

TL;DR

This PR adds a DuckDBQuery task plugin that runs queries using DuckDB as the DBMS.

Type

• Bug Fix
• Feature
• Plugin

Are all requirements met?

• Code completed
• Smoke tested
• Unit tests added
• Code documentation added
• Any pending items have an associated Issue

Complete description

Capturing the crucial assumptions I made while building the task plugin:

• The DuckDBQuery task parameter that a user needs to send argument to includes query and can contemplate adding includes inputs.
• query can include a set of queries that'll be run sequentially. The last query needs to be a SELECT query.
• inputs can include structured dataset or a list of parameters to be sent to the queries.
• The output is a pyarrow table. Can be converted to any structured dataset compatible type.
• The connection mode is set to :memory, i.e., the data is always stored in an in-memory, non-persistent database. It can be set to a file, but it's difficult to make the file accessible to different DuckDBQuery pods, which otherwise wouldn't make sense because file is persistent, and it needs to be leveraged.

Example:

duckdb_query = DuckDBQuery(
    name="read_parquet",
    query=[
        "INSTALL httpfs",
        "LOAD httpfs",
        """SELECT hour(lpep_pickup_datetime) AS hour, count(*) AS count FROM READ_PARQUET(?) GROUP BY hour""",
    ],
    inputs=kwtypes(params=list[str]),
)

@workflow
def parquet_wf(parquet_file: str) -> pd.DataFrame:
    return duckdb_query(params=[parquet_file])

assert isinstance(
    parquet_wf(parquet_file="https://d37ci6vzurychx.cloudfront.net/trip-data/green_tripdata_2022-02.parquet"),
    pd.DataFrame,
)

Tracking Issue

Fixes flyteorg/flyte#3246, flyteorg/flyte#2865

Follow-up issue

NA
OR
https://github.com/flyteorg/flyte/issues/

[codecov]

Codecov Report

Merging #1419 (9a0e42b) into master (b3ad158) will not change coverage.
The diff coverage is n/a.

@@           Coverage Diff           @@
##           master    #1419   +/-   ##
=======================================
  Coverage   69.32%   69.32%           
=======================================
  Files         305      305           
  Lines       28671    28671           
  Branches     2718     2718           
=======================================
  Hits        19877    19877           
  Misses       8276     8276           
  Partials      518      518           

Impacted Files	Coverage Δ
plugins/setup.py	0.00% <ø> (ø)

Help us with your feedback. Take ten seconds to tell us how you rate us. Have a feature suggestion? Share it here.

[kumare3]

@samhita-alla can we support accepting StructuredDataset as an input?

[samhita-alla]

@kumare3, DuckDBQuery task accepts structured dataset already. Here's an example:

duckdb_task = DuckDBQuery(
    name="duckdb_sd_df",
    query="SELECT * FROM pandas_df WHERE i = 2",
    inputs=kwtypes(pandas_df=StructuredDataset),
)

@task
def get_pandas_df() -> StructuredDataset:
    return StructuredDataset(
        dataframe=pd.DataFrame.from_dict({"i": [1, 2, 3, 4], "j": ["one", "two", "three", "four"]})
    )

@workflow
def pandas_wf(pandas_df: StructuredDataset) -> pd.DataFrame:
    return duckdb_task(pandas_df=pandas_df)

assert isinstance(pandas_wf(pandas_df=get_pandas_df()), pd.DataFrame)

Let me know if you're looking for something different.

[cosmicBboy]

looking great! added some comments for docstrings

[samhita-alla]

@cosmicBboy, thanks for reviewing the PR! Can you look through it again?

[cosmicBboy]

DuckDB api reference is blank, I think we need to update the https://github.com/flyteorg/flytekit/blob/master/doc-requirements.in file with duckdb as a dependency, or else it won't be rendered.

[cosmicBboy]

@samhita-alla we'll also need to invest a bit in enable warnings as errors in the sphinx build process -W like we did with flytesnacks. It'll let us catch these kinds of docs build errors

[samhita-alla]

@cosmicBboy, fixed the docs and added a GitHub action to show warnings as errors.

[pingsutw]

Could we add an Output Schema type here like snowflake? if output_schema_type is none, we won't generate output dataset.
And the type should change to StructuredDataset, because we already deprecated FlyteSchema

[samhita-alla]

@pingsutw, this task isn't helpful if there's no output dataset. The DuckDBQuery task runs some queries and returns the output of a SELECT statement. Hence, it must return the query output, and in this case, it's StructuredDataset. Also, I can definitely add output_schema_type. But the output has to always be a StructuredDataset. So is it necessary? I'm already hard-coding the output type in the initialization. Let me know what you think.

[pingsutw]

I see. do we support insert or some other operations? If not, I think we don't need schema type for now

[samhita-alla]

Yeah, we do. So you can give a bunch of queries to a single DuckDBQuery task. But the last one needs to be a SELECT query because after say, you insert the data, you need to retrieve the data, right? Else, it's of no use. I'm using the non-persistent offering by DuckDB. So all the data will be available only within the query. Does that make sense?

[pingsutw]

make sense. Thanks for the explanation.

[samhita-alla]

Thanks, Kevin! Will merge this PR after @cosmicBboy approves as well.

[samhita-alla]

@kumare3, let me know if this PR looks good to you.