Skip to content

Commit 85f1895

Browse files
committed
Fix and conform Spark API's show() and UDF register().
Match DataFrame.show() signature with PySpark's. Translate show() arguments to those of DuckDBPyRelation's show(). Add third argument to create_function() call in register(). Assume register() takes Spark API return type argument. Add null_handling argument to match PySpark's behavior.
1 parent 331c741 commit 85f1895

2 files changed

Lines changed: 27 additions & 3 deletions

File tree

duckdb/experimental/spark/sql/dataframe.py

Lines changed: 20 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -38,8 +38,26 @@ def __init__(self, relation: duckdb.DuckDBPyRelation, session: "SparkSession") -
3838
if self.relation is not None:
3939
self._schema = duckdb_to_spark_schema(self.relation.columns, self.relation.types)
4040

41-
def show(self, **kwargs) -> None: # noqa: D102
42-
self.relation.show()
41+
def show(
42+
self,
43+
n: int = 20,
44+
truncate: Union[bool, int] = True,
45+
vertical: bool = False,
46+
) -> None: # noqa: D102
47+
if isinstance(truncate, int):
48+
max_col_width = truncate
49+
elif truncate:
50+
max_col_width = 20
51+
else:
52+
max_col_width = 9999
53+
54+
render_mode = 'COLUMNS' if vertical else None
55+
56+
self.relation.show(
57+
max_rows=n,
58+
max_col_width=max_col_width,
59+
render_mode=render_mode
60+
)
4361

4462
def toPandas(self) -> "PandasDataFrame": # noqa: D102
4563
return self.relation.df()

duckdb/experimental/spark/sql/udf.py

Lines changed: 7 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -23,7 +23,13 @@ def register( # noqa: D102
2323
f: "Callable[..., Any] | UserDefinedFunctionLike",
2424
returnType: Optional["DataTypeOrString"] = None,
2525
) -> "UserDefinedFunctionLike":
26-
self.sparkSession.conn.create_function(name, f, return_type=returnType)
26+
self.sparkSession.conn.create_function(
27+
name,
28+
f,
29+
None,
30+
return_type=returnType.duckdb_type if returnType else None,
31+
null_handling='special',
32+
)
2733

2834
def registerJavaFunction( # noqa: D102
2935
self,

0 commit comments

Comments
 (0)