Skip to content

Commit 27e4f30

Browse files
committed
Initial commit for creating wrapper classes and functions for all user facing python features
1 parent d00c00a commit 27e4f30

15 files changed

+4068
-78
lines changed

python/datafusion/__init__.py

Lines changed: 53 additions & 53 deletions
Original file line numberDiff line numberDiff line change
@@ -25,64 +25,67 @@
2525

2626
import pyarrow as pa
2727

28-
from ._internal import (
29-
AggregateUDF,
30-
Config,
31-
DataFrame,
28+
from .context import (
3229
SessionContext,
3330
SessionConfig,
3431
RuntimeConfig,
35-
ScalarUDF,
3632
SQLOptions,
3733
)
3834

35+
# The following imports are okay to remain as opaque to the user.
36+
from ._internal import Config
37+
38+
from .udf import ScalarUDF, AggregateUDF
39+
3940
from .common import (
4041
DFSchema,
4142
)
4243

44+
from .dataframe import DataFrame
45+
4346
from .expr import (
44-
Alias,
45-
Analyze,
47+
# Alias,
48+
# Analyze,
4649
Expr,
47-
Filter,
48-
Limit,
49-
Like,
50-
ILike,
51-
Projection,
52-
SimilarTo,
53-
ScalarVariable,
54-
Sort,
55-
TableScan,
56-
Not,
57-
IsNotNull,
58-
IsTrue,
59-
IsFalse,
60-
IsUnknown,
61-
IsNotTrue,
62-
IsNotFalse,
63-
IsNotUnknown,
64-
Negative,
65-
InList,
66-
Exists,
67-
Subquery,
68-
InSubquery,
69-
ScalarSubquery,
70-
GroupingSet,
71-
Placeholder,
72-
Case,
73-
Cast,
74-
TryCast,
75-
Between,
76-
Explain,
77-
CreateMemoryTable,
78-
SubqueryAlias,
79-
Extension,
80-
CreateView,
81-
Distinct,
82-
DropTable,
83-
Repartition,
84-
Partitioning,
85-
Window,
50+
# Filter,
51+
# Limit,
52+
# Like,
53+
# ILike,
54+
# Projection,
55+
# SimilarTo,
56+
# ScalarVariable,
57+
# Sort,
58+
# TableScan,
59+
# Not,
60+
# IsNotNull,
61+
# IsTrue,
62+
# IsFalse,
63+
# IsUnknown,
64+
# IsNotTrue,
65+
# IsNotFalse,
66+
# IsNotUnknown,
67+
# Negative,
68+
# InList,
69+
# Exists,
70+
# Subquery,
71+
# InSubquery,
72+
# ScalarSubquery,
73+
# GroupingSet,
74+
# Placeholder,
75+
# Case,
76+
# Cast,
77+
# TryCast,
78+
# Between,
79+
# Explain,
80+
# CreateMemoryTable,
81+
# SubqueryAlias,
82+
# Extension,
83+
# CreateView,
84+
# Distinct,
85+
# DropTable,
86+
# Repartition,
87+
# Partitioning,
88+
# Window,
8689
WindowFrame,
8790
)
8891

@@ -96,7 +99,6 @@
9699
"SQLOptions",
97100
"RuntimeConfig",
98101
"Expr",
99-
"AggregateUDF",
100102
"ScalarUDF",
101103
"Window",
102104
"WindowFrame",
@@ -175,8 +177,6 @@ def column(value):
175177

176178

177179
def literal(value):
178-
if not isinstance(value, pa.Scalar):
179-
value = pa.scalar(value)
180180
return Expr.literal(value)
181181

182182

@@ -200,20 +200,20 @@ def udf(func, input_types, return_type, volatility, name=None):
200200
)
201201

202202

203-
def udaf(accum, input_type, return_type, state_type, volatility, name=None):
203+
def udaf(accum, input_types, return_type, state_type, volatility, name=None):
204204
"""
205205
Create a new User Defined Aggregate Function
206206
"""
207207
if not issubclass(accum, Accumulator):
208208
raise TypeError("`accum` must implement the abstract base class Accumulator")
209209
if name is None:
210210
name = accum.__qualname__.lower()
211-
if isinstance(input_type, pa.lib.DataType):
212-
input_type = [input_type]
211+
if isinstance(input_types, pa.lib.DataType):
212+
input_types = [input_types]
213213
return AggregateUDF(
214214
name=name,
215215
accumulator=accum,
216-
input_type=input_type,
216+
input_types=input_types,
217217
return_type=return_type,
218218
state_type=state_type,
219219
volatility=volatility,

python/datafusion/catalog.py

Lines changed: 42 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,42 @@
1+
from __future__ import annotations
2+
3+
import datafusion._internal as df_internal
4+
5+
from typing import TYPE_CHECKING
6+
7+
if TYPE_CHECKING:
8+
import pyarrow
9+
10+
11+
class Catalog:
12+
def __init__(self, catalog: df_internal.Catalog) -> None:
13+
self.catalog = catalog
14+
15+
def names(self) -> list[str]:
16+
return self.catalog.names()
17+
18+
def database(self, name: str = "public") -> Database:
19+
return Database(self.catalog.database(name))
20+
21+
22+
class Database:
23+
def __init__(self, db: df_internal.Database) -> None:
24+
self.db = db
25+
26+
def names(self) -> set[str]:
27+
return self.db.names()
28+
29+
def table(self, name: str) -> Table:
30+
return Table(self.db.table(name))
31+
32+
33+
class Table:
34+
def __init__(self, table: df_internal.Table) -> None:
35+
self.table = table
36+
37+
def schema(self) -> pyarrow.Schema:
38+
return self.table.schema()
39+
40+
@property
41+
def kind(self) -> str:
42+
return self.table.kind()

0 commit comments

Comments
 (0)