From f45d856feea5c3e86a410e9f4ab9d3817a3cce3a Mon Sep 17 00:00:00 2001 From: Marc Garcia Date: Wed, 10 Jul 2024 14:44:16 +0700 Subject: [PATCH] Expose the correct namespaces for Python modules --- python/datafusion/__init__.py | 129 ++++++++++++++++-------------- python/datafusion/common.py | 13 ++- python/datafusion/expr.py | 13 ++- python/datafusion/functions.py | 13 ++- python/datafusion/object_store.py | 13 ++- python/datafusion/substrait.py | 13 ++- 6 files changed, 119 insertions(+), 75 deletions(-) diff --git a/python/datafusion/__init__.py b/python/datafusion/__init__.py index 846b1a45..cd30bb53 100644 --- a/python/datafusion/__init__.py +++ b/python/datafusion/__init__.py @@ -88,66 +88,6 @@ __version__ = importlib_metadata.version(__name__) -__all__ = [ - "Config", - "DataFrame", - "SessionContext", - "SessionConfig", - "SQLOptions", - "RuntimeConfig", - "Expr", - "AggregateUDF", - "ScalarUDF", - "Window", - "WindowFrame", - "column", - "literal", - "TableScan", - "Projection", - "DFSchema", - "DFField", - "Analyze", - "Sort", - "Limit", - "Filter", - "Like", - "ILike", - "SimilarTo", - "ScalarVariable", - "Alias", - "Not", - "IsNotNull", - "IsTrue", - "IsFalse", - "IsUnknown", - "IsNotTrue", - "IsNotFalse", - "IsNotUnknown", - "Negative", - "ScalarFunction", - "BuiltinScalarFunction", - "InList", - "Exists", - "Subquery", - "InSubquery", - "ScalarSubquery", - "GroupingSet", - "Placeholder", - "Case", - "Cast", - "TryCast", - "Between", - "Explain", - "SubqueryAlias", - "Extension", - "CreateMemoryTable", - "CreateView", - "Distinct", - "DropTable", - "Repartition", - "Partitioning", -] - class Accumulator(metaclass=ABCMeta): @abstractmethod @@ -175,6 +115,8 @@ def column(value): def literal(value): + import pyarrow as pa + if not isinstance(value, pa.Scalar): value = pa.scalar(value) return Expr.literal(value) @@ -204,6 +146,8 @@ def udaf(accum, input_type, return_type, state_type, volatility, name=None): """ Create a new User Defined Aggregate Function """ + import pyarrow as pa + if not issubclass(accum, Accumulator): raise TypeError("`accum` must implement the abstract base class Accumulator") if name is None: @@ -218,3 +162,68 @@ def udaf(accum, input_type, return_type, state_type, volatility, name=None): state_type=state_type, volatility=volatility, ) + + +del ABCMeta +del abstractmethod +del List +del importlib_metadata +del pa + + +__all__ = [ + "Config", + "DataFrame", + "SessionContext", + "SessionConfig", + "SQLOptions", + "RuntimeConfig", + "Expr", + "AggregateUDF", + "ScalarUDF", + "Window", + "WindowFrame", + "column", + "literal", + "TableScan", + "Projection", + "DFSchema", + "Analyze", + "Sort", + "Limit", + "Filter", + "Like", + "ILike", + "SimilarTo", + "ScalarVariable", + "Alias", + "Not", + "IsNotNull", + "IsTrue", + "IsFalse", + "IsUnknown", + "IsNotTrue", + "IsNotFalse", + "IsNotUnknown", + "Negative", + "InList", + "Exists", + "Subquery", + "InSubquery", + "ScalarSubquery", + "GroupingSet", + "Placeholder", + "Case", + "Cast", + "TryCast", + "Between", + "Explain", + "SubqueryAlias", + "Extension", + "CreateMemoryTable", + "CreateView", + "Distinct", + "DropTable", + "Repartition", + "Partitioning", +] diff --git a/python/datafusion/common.py b/python/datafusion/common.py index dd56640a..c4faa1b5 100644 --- a/python/datafusion/common.py +++ b/python/datafusion/common.py @@ -16,8 +16,15 @@ # under the License. -from ._internal import common - - def __getattr__(name): + from ._internal import common + return getattr(common, name) + + +def __dir__(): + from ._internal import common + + return list(globals().keys()) + [ + obj for obj in dir(common) if not obj.startswith("_") + ] diff --git a/python/datafusion/expr.py b/python/datafusion/expr.py index e914b85d..759788a6 100644 --- a/python/datafusion/expr.py +++ b/python/datafusion/expr.py @@ -16,8 +16,15 @@ # under the License. -from ._internal import expr - - def __getattr__(name): + from ._internal import expr + return getattr(expr, name) + + +def __dir__(): + from ._internal import expr + + return list(globals().keys()) + [ + obj for obj in dir(expr) if not obj.startswith("_") + ] diff --git a/python/datafusion/functions.py b/python/datafusion/functions.py index 782ecba2..bb1753d4 100644 --- a/python/datafusion/functions.py +++ b/python/datafusion/functions.py @@ -16,8 +16,15 @@ # under the License. -from ._internal import functions - - def __getattr__(name): + from ._internal import functions + return getattr(functions, name) + + +def __dir__(): + from ._internal import functions + + return list(globals().keys()) + [ + obj for obj in dir(functions) if not obj.startswith("_") + ] diff --git a/python/datafusion/object_store.py b/python/datafusion/object_store.py index 70ecbd2b..54ac9681 100644 --- a/python/datafusion/object_store.py +++ b/python/datafusion/object_store.py @@ -16,8 +16,15 @@ # under the License. -from ._internal import object_store - - def __getattr__(name): + from ._internal import object_store + return getattr(object_store, name) + + +def __dir__(): + from ._internal import object_store + + return list(globals().keys()) + [ + obj for obj in dir(object_store) if not obj.startswith("_") + ] diff --git a/python/datafusion/substrait.py b/python/datafusion/substrait.py index eff809a0..e66f6fe8 100644 --- a/python/datafusion/substrait.py +++ b/python/datafusion/substrait.py @@ -16,8 +16,15 @@ # under the License. -from ._internal import substrait - - def __getattr__(name): + from ._internal import substrait + return getattr(substrait, name) + + +def __dir__(): + from ._internal import substrait + + return list(globals().keys()) + [ + obj for obj in dir(substrait) if not obj.startswith("_") + ]