[perf](udf) Write terminal UDF output directly to Parquet - #541
kaka11chen wants to merge 1 commit into
Conversation
![chatgpt-codex-connector[bot]](https://avatars.githubusercontent.com/in/1144995?s=60&v=4){kind=small}
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 9425e9857b
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
| if "://" in directory or directory.startswith("file:"): | ||
| filesystem, filesystem_directory = pafs.FileSystem.from_uri(directory) | ||
| file_path = posixpath.join(filesystem_directory.rstrip("/"), file_name) |
There was a problem hiding this comment.
When a terminal UDF writes to a remote URI using credentials configured through DuckDB (for example, an S3 CREATE SECRET), constructing a new PyArrow filesystem with FileSystem.from_uri() bypasses DuckDB's filesystem instance and its secret/configuration state. The same write_parquet call that previously used DuckDB's configured filesystem will therefore fail unless equivalent credentials happen to be available independently to PyArrow on every worker; the direct sink must propagate the configured filesystem credentials or decline fusion for these destinations.
Useful? React with 👍 / 👎.
| def _expected_schema(payload: dict[str, Any]) -> pa.Schema: | ||
| names = [str(value) for value in payload.get(_EXPECTED_NAMES) or []] | ||
| type_names = [str(value) for value in payload.get(_EXPECTED_TYPES) or []] | ||
| if not names or len(names) != len(type_names): | ||
| raise ValueError("terminal Arrow Parquet sink requires matching expected names and types") | ||
| return pa.schema([pa.field(name, _arrow_type_from_name(type_name)) for name, type_name in zip(names, type_names)]) |
There was a problem hiding this comment.
For a terminal UDF whose declared schema contains a valid DuckDB type not handled by _arrow_type_from_duckdb_pytype, such as UUID, TIMESTAMPTZ, or UHUGEINT, this unconditional schema reconstruction raises unsupported DuckDB type id before any output is written. Nonempty outputs of those types previously flowed through DuckDB's Arrow conversion and Parquet writer, so either map their Arrow representations here or avoid the fused path for unsupported schemas.
Useful? React with 👍 / 👎.
| options.emplace_back("row_group_size", Value::UBIGINT(bind_data.row_group_size)); | ||
| options.emplace_back("dictionary_pagesize_limit", Value::UBIGINT(bind_data.string_dictionary_page_size_limit)); | ||
| options.emplace_back("version", Value(bind_data.parquet_version == ParquetVersion::V2 ? "2.6" : "1.0")); | ||
| options.emplace_back("data_page_version", Value("1.0")); |
There was a problem hiding this comment.
When the terminal UDF schema includes TIMESTAMP_NS, the default DuckDB ParquetVersion::V1 is translated to PyArrow version="1.0"; PyArrow cannot represent nanosecond timestamps in that format version without coercion and raises when values would lose precision. DuckDB's V1/V2 setting controls permitted encodings rather than PyArrow's logical-type format version, so this mapping makes a previously writable default-schema case fail; use a PyArrow format version that supports nanoseconds while configuring encodings separately.
Useful? React with 👍 / 👎.
Summary
COPYand stream its Arrow output directly into a PyArrowParquetWriter, avoiding the Arrow-to-DuckDB-to-Parquet round trip.Related issue
N/A
Documentation impact
AstroVela/vane-website.The public
Relation.write_parquetAPI is unchanged. The intentional behavior change for unsupported terminal-UDF COPY shapes/options is documented in the summary above.Validation
SKBUILD_BUILD_DIR="$PWD/build/python-release" SKBUILD_CMAKE_BUILD_TYPE=Release uv pip install . --no-build-isolation.VCPKG_INSTALLED_DIR="$PWD/vcpkg_installed" scripts/run_native_tests.sh "[distributed]"— 187 test cases and 7,988 assertions passed.scripts/run_release_tests.sh— 538 passed (15 deselected), plus 11 real-Ray tests passed (542 deselected).scripts/run_fast_tests.sh— 6,206 passed, 602 skipped, 120 deselected, 7 xfailed, and 1 xpassed in the non-Ray shard; 103 passed and 6 skipped in the shared-Ray shard; all seven CPU-only Ray-owner tests passed, with the unavailable vLLM test skipped.scripts/format workspace --changed --check, pre-commit on every changed file, andgit diff --checkpassed.batch_size=64,row_group_size=32) and 17-row Ray input (batch_size=4,row_group_size=3) across task and actor backends.Checklist
licenses and are recorded where required.
weights, or build artifacts are included.
and untrusted input have been considered.
changes passed relevant checks.