Skip to content

Obtain filepath from Dask DataFrame #1145

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
ayushdg merged 8 commits into from
May 26, 2023
Merged

Obtain filepath from Dask DataFrame #1145

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
8 commits
Select commit Hold shift + click to select a range
69c1c22
use hlg_layer creation_info
sarahyurick May 12, 2023
0495211
add keyerror
sarahyurick May 12, 2023
b0a245e
style fix
sarahyurick May 12, 2023
92701e4
use is_dataframe_like
sarahyurick May 12, 2023
86d8f3b
add hasattr
sarahyurick May 12, 2023
bcf74db
use parquet_ddf fixture
sarahyurick May 19, 2023
bf10f85
Merge branch 'main' into read_parquet_filepath
sarahyurick May 24, 2023
b35ea00
Merge branch 'main' into read_parquet_filepath
ayushdg May 25, 2023
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
10 changes: 10 additions & 0 deletions dask_sql/context.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -8,6 +8,7 @@
import pandas as pd
from dask import config as dask_config
from dask.base import optimize
from dask.utils_test import hlg_layer

from dask_planner.rust import (
DaskSchema,
Expand Down Expand Up @@ -247,6 +248,15 @@ def create_table(
if type(input_table) == str:
dc.filepath = input_table
self.schema[schema_name].filepaths[table_name.lower()] = input_table
elif hasattr(input_table, "dask") and dd.utils.is_dataframe_like(input_table):
try:
dask_filepath = hlg_layer(
input_table.dask, "read-parquet"
).creation_info["args"][0]
dc.filepath = dask_filepath
self.schema[schema_name].filepaths[table_name.lower()] = dask_filepath
except KeyError:
logger.debug("Expected 'read-parquet' layer")

if parquet_statistics and not statistics:
statistics = parquet_statistics(dc.df)
Expand Down
37 changes: 17 additions & 20 deletions tests/unit/test_context.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -312,33 +312,30 @@ def test_alter_table(c, df_simple):
del c.schema[c.schema_name].tables["physics"]


def test_filepath(tmpdir):
def test_filepath(tmpdir, parquet_ddf):
c = Context()

parquet_path = os.path.join(tmpdir, "parquet")
parquet_df = pd.DataFrame(
{
"a": [1, 2, 3] * 5,
"b": range(15),
"c": ["A"] * 15,
"d": [
pd.Timestamp("2013-08-01 23:00:00"),
pd.Timestamp("2014-09-01 23:00:00"),
pd.Timestamp("2015-10-01 23:00:00"),
]
* 5,
"index": range(15),
},
)
dd.from_pandas(parquet_df, npartitions=3).to_parquet(parquet_path)
c.create_table("parquet_df", parquet_path, format="parquet")

assert c.schema["root"].tables["parquet_df"].filepath == parquet_path
assert c.schema["root"].filepaths["parquet_df"] == parquet_path
# Create table with string (Parquet filepath)
c.create_table("parquet_ddf", parquet_path, format="parquet")

assert c.schema["root"].tables["parquet_ddf"].filepath == parquet_path
assert c.schema["root"].filepaths["parquet_ddf"] == parquet_path

df = pd.DataFrame({"a": [2, 1, 2, 3], "b": [3, 3, 1, 3]})
c.create_table("df", df)

assert c.schema["root"].tables["df"].filepath is None
with pytest.raises(KeyError):
c.schema["root"].filepaths["df"]


def test_ddf_filepath(tmpdir, parquet_ddf):
c = Context()
parquet_path = os.path.join(tmpdir, "parquet")

# Create table with Dask DataFrame (created from read_parquet)
c.create_table("parquet_ddf", parquet_ddf)

assert c.schema["root"].tables["parquet_ddf"].filepath == parquet_path
assert c.schema["root"].filepaths["parquet_ddf"] == parquet_path