def test_query(db):
import sqlalchemy as sa
from sqlalchemy import sql
s1 = sql.select([sql.column("number"),
sql.column("name")]).select_from(sql.table("test"))
out = read_sql_query(s1, db, npartitions=2, index_col="number")
assert_eq(out, df[["name"]])
s2 = (sql.select([
sa.cast(sql.column("number"), sa.types.BigInteger).label("number"),
sql.column("name"),
]).where(sql.column("number") >= 5).select_from(sql.table("test")))
out = read_sql_query(s2, db, npartitions=2, index_col="number")
assert_eq(out, df.loc[5:, ["name"]])
def test_query_index_from_query(db):
from sqlalchemy import sql
number = sql.column("number")
name = sql.column("name")
s1 = sql.select([number, name, sql.func.length(name).label("lenname")]).select_from(
sql.table("test")
)
out = read_sql_query(s1, db, npartitions=2, index_col="lenname")
lenname_df = df.copy()
lenname_df["lenname"] = lenname_df["name"].str.len()
lenname_df = lenname_df.reset_index().set_index("lenname")
assert_eq(out, lenname_df.loc[:, ["number", "name"]])
def test_query_with_meta(db):
from sqlalchemy import sql
data = {
"name": pd.Series([], name="name", dtype="str"),
"age": pd.Series([], name="age", dtype="int"),
}
index = pd.Index([], name="number", dtype="int")
meta = pd.DataFrame(data, index=index)
s1 = sql.select(
[sql.column("number"), sql.column("name"), sql.column("age")]
).select_from(sql.table("test"))
out = read_sql_query(s1, db, npartitions=2, index_col="number", meta=meta)
# Don't check dtype for windows https://github.com/dask/dask/issues/8620
assert_eq(out, df[["name", "age"]], check_dtype=sys.platform != "win32")