Created
July 29, 2025 01:53
-
-
Save ykdojo/3e1d5df460c1e68715352b968f16394e to your computer and use it in GitHub Desktop.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| #!/usr/bin/env python3 | |
| """Test ellipsis syntax by executing UDFs and checking resulting column types""" | |
| import daft | |
| from daft import col | |
| # Create test dataframe | |
| df = daft.from_pydict({"values": [1.0, 2.0, 3.0]}) | |
| print("=== Testing list[type, ...] Ellipsis Syntax by Execution ===\n") | |
| # Test 1: Simple list with ellipsis | |
| print("1. Simple list with ellipsis:") | |
| @daft.udf(return_dtype=list[float, ...]) | |
| def simple_ellipsis(x): | |
| return [[v * 2.0] for v in x.to_pylist()] | |
| result1 = df.with_column("simple_result", simple_ellipsis(col("values"))) | |
| print(f" list[float, ...] → {result1.schema()['simple_result'].dtype}") | |
| print() | |
| # Test 2: Nested lists with ellipsis | |
| print("2. Nested lists with ellipsis:") | |
| @daft.udf(return_dtype=list[list[int, ...], ...]) | |
| def nested_ellipsis(x): | |
| return [[[int(v), int(v)+1]] for v in x.to_pylist()] | |
| result2 = df.with_column("nested_result", nested_ellipsis(col("values"))) | |
| print(f" list[list[int, ...], ...] → {result2.schema()['nested_result'].dtype}") | |
| print() | |
| # Test 3: Complex type with ellipsis - list of list of struct of list | |
| print("3. Complex nested type with ellipsis:") | |
| @daft.udf(return_dtype=list[list[{"foo": list[float, ...]}, ...], ...]) | |
| def complex_ellipsis(x): | |
| # Return the correct nested structure | |
| return [ | |
| [ | |
| [{"foo": [v, v * 2.0]}] | |
| ] | |
| for v in x.to_pylist() | |
| ] | |
| result3 = df.with_column("complex_result", complex_ellipsis(col("values"))) | |
| dtype3 = result3.schema()['complex_result'].dtype | |
| print(f" list[list[{{\"foo\": list[float, ...]}}, ...], ...] → {dtype3}") | |
| print(f" Struct key 'foo' present? {'foo' in str(dtype3)}") | |
| print() | |
| # Test 4: The proposed syntax from issue #4829 | |
| print("4. Testing the proposed [type, ...] syntax:") | |
| try: | |
| @daft.udf(return_dtype=[float, ...]) | |
| def proposed_syntax(x): | |
| return [[v * 3.0] for v in x.to_pylist()] | |
| result4 = df.with_column("proposed", proposed_syntax(col("values"))) | |
| print(f" [float, ...] → {result4.schema()['proposed'].dtype}") | |
| except Exception as e: | |
| print(f" [float, ...] → ❌ {type(e).__name__}") | |
| print("\n=== Key Findings ===") | |
| print("✅ list[type, ...] works - standard Python syntax") | |
| print("✅ Complex nested types work with ellipsis") | |
| print("❌ [type, ...] doesn't work - not valid Python") | |
| """ | |
| Tested with: | |
| - Python version: 3.12.10 | |
| - Daft version: 0.5.13 | |
| Actual output from test run: | |
| 1. Simple list with ellipsis: | |
| list[float, ...] → List[Float64] | |
| 2. Nested lists with ellipsis: | |
| list[list[int, ...], ...] → List[List[Int64]] | |
| 3. Complex nested type with ellipsis: | |
| list[list[{"foo": list[float, ...]}, ...], ...] → List[List[Struct[foo: List[Float64]]]] | |
| Struct key 'foo' present? True | |
| 4. Testing the proposed [type, ...] syntax: | |
| [float, ...] → ❌ ValueError | |
| """ |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment