JunRuiLee commented on code in PR #420: URL: https://github.com/apache/paimon-rust/pull/420#discussion_r3497489071
########## bindings/python/src/write.rs: ########## @@ -0,0 +1,203 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +use std::sync::Arc; + +use arrow::datatypes::Schema as ArrowSchema; +use arrow::pyarrow::FromPyArrow; +use arrow::record_batch::RecordBatch; +use paimon::table::{CommitMessage, Table, TableCommit, TableWrite}; +use paimon_datafusion::runtime::runtime; +use pyo3::exceptions::{PyTypeError, PyValueError}; +use pyo3::prelude::*; + +use crate::error::to_py_err; + +/// Validate an incoming batch schema against the table's target Arrow schema: +/// field count, order, and names must match, and types must match exactly. The +/// nullable flag is intentionally NOT compared, since `build_target_arrow_schema` +/// derives nullability from the Paimon field while pyarrow-constructed batches +/// infer nullable=true. No cast — callers supply correctly-typed batches. +/// +/// Type matching is strict (no binary-family interchange): the lower write path +/// downcasts to the exact Arrow array for each Paimon type (e.g. a `Binary` / +/// `VarBinary` field requires `arrow_array::BinaryArray`, not `LargeBinary` / +/// `FixedSizeBinary`). Accepting a near-equivalent type here would pass +/// validation but then fail deeper with a type-mismatch (or write files whose +/// Arrow schema differs from the table), so it is rejected up front. +fn validate_batch_schema(input: &ArrowSchema, target: &ArrowSchema) -> PyResult<()> { + let mismatch = || { + PyValueError::new_err(format!( + "Input schema is not consistent with the table schema. \ + input: {input:?}, table: {target:?}" + )) + }; + if input.fields().len() != target.fields().len() { + return Err(mismatch()); + } + for (i, t) in input.fields().iter().zip(target.fields().iter()) { + if i.name() != t.name() { + return Err(mismatch()); + } + if i.data_type() != t.data_type() { + return Err(mismatch()); + } + } + Ok(()) +} + +/// Builder for the batch write loop, created via [`crate::table::PyTable::new_write_builder`]. +/// +/// Holds the owning table plus a single fixed `commit_user`, generated once and +/// shared by both `new_write()` and `new_commit()` so that writers and the +/// committer agree on the commit user (Paimon uses it for duplicate-commit +/// detection). Creating a fresh `WriteBuilder` per call would otherwise mint a +/// new random UUID each time. +#[pyclass(name = "WriteBuilder", module = "pypaimon_rust.datafusion")] +pub struct PyWriteBuilder { + table: Arc<Table>, + commit_user: String, +} + +impl PyWriteBuilder { + pub fn new(table: Arc<Table>) -> Self { + let commit_user = table.new_write_builder().commit_user().to_string(); + Self { table, commit_user } + } +} + +#[pymethods] +impl PyWriteBuilder { + /// Create a writer for accumulating Arrow batches. + fn new_write(&self) -> PyResult<PyTableWrite> { + let builder = self + .table + .new_write_builder() + .with_commit_user(self.commit_user.clone()) + .map_err(to_py_err)?; + let target_schema = paimon::arrow::build_target_arrow_schema(self.table.schema().fields()) + .map_err(to_py_err)?; + Ok(PyTableWrite { + inner: builder.new_write().map_err(to_py_err)?, + target_schema, + table_location: self.table.location().to_string(), + }) + } + + /// Create a committer for persisting prepared commit messages. + fn new_commit(&self) -> PyResult<PyTableCommit> { + let builder = self + .table + .new_write_builder() + .with_commit_user(self.commit_user.clone()) + .map_err(to_py_err)?; + Ok(PyTableCommit { + inner: builder.new_commit(), + table_location: self.table.location().to_string(), + }) + } +} + +/// A stateful writer that accumulates Arrow batches until `prepare_commit`. +/// +/// Marked `unsendable`: the underlying `TableWrite` holds file writers that are +/// not `Sync`, so the object enforces single-thread access at runtime. +#[pyclass(name = "TableWrite", module = "pypaimon_rust.datafusion", unsendable)] +pub struct PyTableWrite { + inner: TableWrite, + /// The table's target Arrow schema, used to validate incoming batches. + target_schema: Arc<ArrowSchema>, + /// The owning table's location, stamped onto produced commit messages so a + /// committer can reject messages prepared for a different table. + table_location: String, +} + +#[pymethods] +impl PyTableWrite { + /// Write a single PyArrow RecordBatch into the table's writers. + fn write_arrow(&mut self, py: Python<'_>, batch: &Bound<'_, PyAny>) -> PyResult<()> { + let batch = RecordBatch::from_pyarrow_bound(batch)?; + validate_batch_schema(&batch.schema(), &self.target_schema)?; + let rt = runtime(); + py.detach(|| rt.block_on(async { self.inner.write_arrow_batch(&batch).await })) + .map_err(to_py_err) + } + + /// Close writers and return the commit messages (opaque; pass to commit()). + fn prepare_commit(&mut self, py: Python<'_>) -> PyResult<Vec<PyCommitMessage>> { + let rt = runtime(); + let messages = py + .detach(|| rt.block_on(async { self.inner.prepare_commit().await })) + .map_err(to_py_err)?; + Ok(messages + .into_iter() + .map(|inner| PyCommitMessage { Review Comment: You're right — the location check alone still let a different WriteBuilder (same table) through, and since each builder mints its own commit_user, the writer and committer would disagree on it. Now PyCommitMessage also carries the originating builder's commit_user, and commit() rejects a mismatch in addition to the table-location check. Added a same-table-different-builder regression test (the legit table.new_write_builder() write+commit flow still passes since it shares one builder). Fixed in [6c1337a](https://github.com/apache/paimon-rust/pull/420/commits/6c1337ac51e245839bfdab234bd1809d480b14b0). -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected]
