diff --git a/.github/workflows/python_test.yaml b/.github/workflows/python_test.yaml index 905b9c56d6161..01a36af870af4 100644 --- a/.github/workflows/python_test.yaml +++ b/.github/workflows/python_test.yaml @@ -50,7 +50,7 @@ jobs: run: | source venv/bin/activate flake8 python --ignore=E501 - black --line-length 79 --check python + black --line-length 79 --diff --check python - name: Run tests run: | source venv/bin/activate diff --git a/docs/source/python/index.rst b/docs/source/python/index.rst index 56f9097ffdbdc..57ab8d16bf2e1 100644 --- a/docs/source/python/index.rst +++ b/docs/source/python/index.rst @@ -39,11 +39,10 @@ Simple usage: .. code-block:: python import datafusion + from datafusion import functions as f + from datafusion import col import pyarrow - # an alias - f = datafusion.functions - # create a context ctx = datafusion.ExecutionContext() @@ -56,8 +55,8 @@ Simple usage: # create a new statement df = df.select( - f.col("a") + f.col("b"), - f.col("a") - f.col("b"), + col("a") + col("b"), + col("a") - col("b"), ) # execute and collect the first (and only) batch @@ -77,7 +76,7 @@ UDFs udf = f.udf(is_null, [pyarrow.int64()], pyarrow.bool_()) - df = df.select(udf(f.col("a"))) + df = df.select(udf(col("a"))) UDAF @@ -117,7 +116,7 @@ UDAF df = df.aggregate( [], - [udaf(f.col("a"))] + [udaf(col("a"))] ) diff --git a/python/datafusion/__init__.py b/python/datafusion/__init__.py index 4f9082e7e4022..0a25592f80ae2 100644 --- a/python/datafusion/__init__.py +++ b/python/datafusion/__init__.py @@ -28,6 +28,7 @@ ScalarUDF, ) + __all__ = [ "DataFrame", "ExecutionContext", @@ -61,12 +62,18 @@ def column(value): return Expression.column(value) +col = column + + def literal(value): if not isinstance(value, pa.Scalar): value = pa.scalar(value) return Expression.literal(value) +lit = literal + + def udf(func, input_types, return_type, volatility, name=None): """ Create a new User Defined Function diff --git a/python/src/expression.rs b/python/src/expression.rs index 21cecaa1ccce7..5e1cad246bf87 100644 --- a/python/src/expression.rs +++ b/python/src/expression.rs @@ -90,6 +90,10 @@ impl PyObjectProtocol for PyExpr { }; expr.into() } + + fn __str__(&self) -> PyResult { + Ok(format!("{}", self.expr)) + } } #[pymethods]