8000 Improve col class access using __getattr__ · chenkovsky/datafusion-python@00dea11 · GitHub
[go: up one dir, main page]

Skip to content

Commit 00dea11

Browse files
deanm0000timsaucer
andauthored
Improve col class access using __getattr__
Co-authored-by: Tim Saucer <timsaucer@gmail.com>
1 parent 91b6635 commit 00dea11

File tree

3 files changed

+70
-10
lines changed

3 files changed

+70
-10
lines changed

python/datafusion/__init__.py

Lines changed: 2 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -26,6 +26,8 @@
2626
except ImportError:
2727
import importlib_metadata
2828

29+
from datafusion.col import col, column
30+
2931
from . import functions, object_store, substrait, unparser
3032

3133
# The following imports are okay to remain as opaque to the user.
@@ -95,16 +97,6 @@
9597
]
9698

9799

98-
def column(value: str) -> Expr:
99-
"""Create a column expression."""
100-
return Expr.column(value)
101-
102-
103-
def col(value: str) -> Expr:
104-
"""Create a column expression."""
105-
return Expr.column(value)
106-
107-
108100
def literal(value) -> Expr:
109101
"""Create a literal expression."""
110102
return Expr.literal(value)

python/datafusion/col.py

Lines changed: 45 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,45 @@
1+
# Licensed to the Apache Software Foundation (ASF) under one
2+
# or more contributor license agreements. See the NOTICE file
3+
# distributed with this work for additional information
4+
# regarding copyright ownership. The ASF licenses this file
5+
# to you under the Apache License, Version 2.0 (the
6+
# "License"); you may not use this file except in compliance
7+
# with the License. You may obtain a copy of the License at
8+
#
9+
# http://www.apache.org/licenses/LICENSE-2.0
10+
#
11+
# Unless required by applicable law or agreed to in writing,
12+
# software distributed under the License is distributed on an
13+
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
14+
# KIND, either express or implied. See the License for the
15+
# specific language governing permissions and limitations
16+
# under the License.
17+
18+
"""Col class."""
19+
20+
from datafusion.expr import Expr
21+
22+
23+
class Col:
24+
"""Create a column expression.
25+
26+
This helper class allows an extra syntax of creating columns using the __getattr__
27+
method.
28+
"""
29+
30+
def __call__(self, value: str) -> Expr:
31+
"""Create a column expression."""
32+
return Expr.column(value)
33+
34+
def __getattr__(self, value: str) -> Expr:
35+
"""Create a column using attribute syntax."""
36+
# For autocomplete to work with IPython
37+
if value.startswith("__wrapped__"):
38+
return getattr(type(self), value)
39+
40+
return Expr.column(value)
41+
42+
43+
col: Col = Col()
44+
column: Col = Col()
45+
__all__ = ["col", "column"]

python/tests/test_expr.py

Lines changed: 23 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -249,6 +249,29 @@ def test_fill_null(df):
249249
assert result.column(2) == pa.array([1234, 1234, 8])
250250

251251

252+
def test_col_getattr():
253+
ctx = SessionContext()
254+
data = {
255+
"array_values": [[1, 2, 3], [4, 5], [6], []],
256+
"struct_values": [
257+
{"name": "Alice", "age": 15},
258+
{"name": "Bob", "age": 14},
259+
{"name": "Charlie", "age": 13},
260+
{"name": None, "age": 12},
261+
],
262+
}
263+
df = ctx.from_pydict(data, name="table1")
264+
265+
names = df.select(col.struct_values["name"].alias("name")).collect()
266+
names = [r.as_py() for rs in names for r in rs["name"]]
267+
268+
array_values = df.select(col.array_values[1].alias("value")).collect()
269+
array_values = [r.as_py() for rs in array_values for r in rs["value"]]
270+
271+
assert names == ["Alice", "Bob", "Charlie", None]
272+
assert array_values == [2, 5, None, None]
273+
274+
252275
def test_alias_with_metadata(df):
253276
df = df.select(col("a").alias("b", {"key": "value"}))
254277
assert df.schema().field("b").metadata == {b"key": b"value"}

0 commit comments

Comments
 (0)
0