Google Git
Sign in
apache/fluss-rust/refs/heads/main/./bindings/python/src/table.rs
blob: cbdcbc0134c0a53b98b35e162451521378a9d4e3 [file]
// 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 crate::TOKIO_RUNTIME;
use crate::*;
use arrow::array::RecordBatch as ArrowRecordBatch;
use arrow::record_batch::RecordBatchReader as _;
use arrow_pyarrow::{FromPyArrow, ToPyArrow};
use arrow_schema::SchemaRef;
use fcore::client::LimitBatchScanner;
use fcore::metadata::{DataField, DataType, MapType, RowType};
use fcore::row::binary_array::{FlussArray, FlussArrayWriter};
use fcore::row::binary_map::{FlussMap, FlussMapWriter};
use fcore::row::{Datum, F32, F64, GenericRow, InternalRow};
use fluss::record::to_arrow_schema;
use indexmap::IndexMap;
use pyo3::IntoPyObjectExt;
use pyo3::exceptions::{PyIndexError, PyRuntimeError, PyTypeError};
use pyo3::sync::PyOnceLock;
use pyo3::types::{
IntoPyDict, PyBool, PyByteArray, PyBytes, PyDate, PyDateAccess, PyDateTime, PyDelta,
PyDeltaAccess, PyDict, PyList, PySequence, PySlice, PyString, PyTime, PyTimeAccess, PyTuple,
PyType, PyTzInfo,
};
use pyo3_async_runtimes::tokio::future_into_py;
use std::collections::HashMap;
use std::sync::Arc;
use std::time::Duration;
use tokio::sync::Mutex;
// Time conversion constants
const MILLIS_PER_SECOND: i64 = 1_000;
const MILLIS_PER_MINUTE: i64 = 60_000;
const MILLIS_PER_HOUR: i64 = 3_600_000;
const MICROS_PER_MILLI: i64 = 1_000;
const MICROS_PER_SECOND: i64 = 1_000_000;
const MICROS_PER_DAY: i64 = 86_400_000_000;
const NANOS_PER_MILLI: i64 = 1_000_000;
const NANOS_PER_MICRO: i64 = 1_000;
const DEFAULT_POLL_INTERVAL_MS: i64 = 1000;
/// Represents a single scan record with metadata.
///
/// Matches Rust/Java: offset, timestamp, change_type, row.
/// The bucket is the key in ScanRecords, not on the individual record.
#[pyclass]
pub struct ScanRecord {
#[pyo3(get)]
offset: i64,
#[pyo3(get)]
timestamp: i64,
#[pyo3(get)]
change_type: ChangeType,
/// Store row as a Python dict directly
row_dict: Py<PyDict>,
}
#[pymethods]
impl ScanRecord {
/// Get the row data as a dictionary
#[getter]
pub fn row(&self, py: Python) -> Py<PyDict> {
self.row_dict.clone_ref(py)
}
fn __str__(&self) -> String {
format!(
"ScanRecord(offset={}, timestamp={}, change_type={})",
self.offset,
self.timestamp,
self.change_type.short_string()
)
}
fn __repr__(&self) -> String {
self.__str__()
}
}
impl ScanRecord {
/// Create a ScanRecord from core types
pub fn from_core(
py: Python,
record: &fcore::record::ScanRecord,
row_type: &fcore::metadata::RowType,
) -> PyResult<Self> {
let fields = row_type.fields();
let row = record.row();
let dict = PyDict::new(py);
for (pos, field) in fields.iter().enumerate() {
let value = datum_to_python_value(py, row, pos, field.data_type())?;
dict.set_item(field.name(), value)?;
}
Ok(ScanRecord {
offset: record.offset(),
timestamp: record.timestamp(),
change_type: ChangeType::from_core(*record.change_type()),
row_dict: dict.unbind(),
})
}
}
/// Represents a batch of records with metadata
#[pyclass]
pub struct RecordBatch {
batch: Arc<ArrowRecordBatch>,
#[pyo3(get)]
bucket: TableBucket,
#[pyo3(get)]
base_offset: i64,
#[pyo3(get)]
last_offset: i64,
}
#[pymethods]
impl RecordBatch {
/// Get the Arrow RecordBatch as PyArrow RecordBatch
#[getter]
pub fn batch(&self, py: Python) -> PyResult<Py<PyAny>> {
let pyarrow_batch = self
.batch
.as_ref()
.to_pyarrow(py)
.map_err(|e| FlussError::new_err(format!("Failed to convert batch: {e}")))?;
Ok(pyarrow_batch.unbind())
}
fn __str__(&self) -> String {
format!(
"RecordBatch(bucket={}, base_offset={}, last_offset={}, rows={})",
self.bucket.__str__(),
self.base_offset,
self.last_offset,
self.batch.num_rows()
)
}
fn __repr__(&self) -> String {
self.__str__()
}
}
impl RecordBatch {
/// Create a RecordBatch from core ScanBatch
pub fn from_scan_batch(scan_batch: fcore::record::ScanBatch) -> Self {
RecordBatch {
bucket: TableBucket::from_core(scan_batch.bucket().clone()),
base_offset: scan_batch.base_offset(),
last_offset: scan_batch.last_offset(),
batch: Arc::new(scan_batch.into_batch()),
}
}
}
/// A collection of scan records grouped by bucket.
///
/// Returned by `LogScanner.poll()`. Records are grouped by `TableBucket`.
#[pyclass]
pub struct ScanRecords {
records_by_bucket: IndexMap<TableBucket, Vec<Py<ScanRecord>>>,
total_count: usize,
}
#[pymethods]
impl ScanRecords {
/// List of distinct buckets that have records in this result.
pub fn buckets(&self) -> Vec<TableBucket> {
self.records_by_bucket.keys().cloned().collect()
}
/// Get records for a specific bucket.
///
/// Returns an empty list if the bucket is not present (matches Rust/Java behavior).
pub fn records(&self, py: Python, bucket: &TableBucket) -> Vec<Py<ScanRecord>> {
self.records_by_bucket
.get(bucket)
.map(|recs| recs.iter().map(|r| r.clone_ref(py)).collect())
.unwrap_or_default()
}
/// Total number of records across all buckets.
pub fn count(&self) -> usize {
self.total_count
}
/// Whether the result set is empty.
pub fn is_empty(&self) -> bool {
self.total_count == 0
}
fn __len__(&self) -> usize {
self.total_count
}
/// Type-dispatched indexing:
/// records[0] → ScanRecord (flat index)
/// records[-1] → ScanRecord (negative index)
/// records[1:3] → list[ScanRecord] (slice)
/// records[bucket] → list[ScanRecord] (by bucket)
fn __getitem__(&self, py: Python, key: &Bound<'_, PyAny>) -> PyResult<Py<PyAny>> {
// Try integer index first
if let Ok(mut idx) = key.extract::<isize>() {
let len = self.total_count as isize;
if idx < 0 {
idx += len;
}
if idx < 0 || idx >= len {
return Err(PyIndexError::new_err(format!(
"index {idx} out of range for ScanRecords of size {len}"
)));
}
let idx = idx as usize;
let mut offset = 0;
for recs in self.records_by_bucket.values() {
if idx < offset + recs.len() {
return Ok(recs[idx - offset].clone_ref(py).into_any());
}
offset += recs.len();
}
return Err(PyRuntimeError::new_err(
"internal error: total_count out of sync with records",
));
}
// Try slice
if let Ok(slice) = key.downcast::<PySlice>() {
let indices = slice.indices(self.total_count as isize)?;
let mut result: Vec<Py<ScanRecord>> = Vec::new();
let mut i = indices.start;
while (indices.step > 0 && i < indices.stop) || (indices.step < 0 && i > indices.stop) {
let idx = i as usize;
let mut offset = 0;
for recs in self.records_by_bucket.values() {
if idx < offset + recs.len() {
result.push(recs[idx - offset].clone_ref(py));
break;
}
offset += recs.len();
}
i += indices.step;
}
return Ok(result.into_pyobject(py).unwrap().into_any().unbind());
}
// Try TableBucket
if let Ok(bucket) = key.extract::<TableBucket>() {
let recs = self.records(py, &bucket);
return Ok(recs.into_pyobject(py).unwrap().into_any().unbind());
}
Err(PyTypeError::new_err(
"index must be int, slice, or TableBucket",
))
}
/// Support `bucket in records`.
fn __contains__(&self, bucket: &TableBucket) -> bool {
self.records_by_bucket.contains_key(bucket)
}
/// Mapping protocol: alias for `buckets()`.
pub fn keys(&self) -> Vec<TableBucket> {
self.buckets()
}
/// Mapping protocol: lazy iterator over record lists, one per bucket.
pub fn values(slf: Bound<'_, Self>) -> ScanRecordsBucketIter {
let this = slf.borrow();
let bucket_keys: Vec<TableBucket> = this.records_by_bucket.keys().cloned().collect();
drop(this);
ScanRecordsBucketIter {
owner: slf.unbind(),
bucket_keys,
bucket_idx: 0,
with_keys: false,
}
}
/// Mapping protocol: lazy iterator over `(TableBucket, list[ScanRecord])` pairs.
pub fn items(slf: Bound<'_, Self>) -> ScanRecordsBucketIter {
let this = slf.borrow();
let bucket_keys: Vec<TableBucket> = this.records_by_bucket.keys().cloned().collect();
drop(this);
ScanRecordsBucketIter {
owner: slf.unbind(),
bucket_keys,
bucket_idx: 0,
with_keys: true,
}
}
fn __str__(&self) -> String {
format!(
"ScanRecords(records={}, buckets={})",
self.total_count,
self.records_by_bucket.len()
)
}
fn __repr__(&self) -> String {
self.__str__()
}
/// Flat iterator over all records across all buckets (matches Java/Rust).
fn __iter__(slf: Bound<'_, Self>) -> ScanRecordsIter {
let this = slf.borrow();
let bucket_keys: Vec<TableBucket> = this.records_by_bucket.keys().cloned().collect();
drop(this);
ScanRecordsIter {
owner: slf.unbind(),
bucket_keys,
bucket_idx: 0,
rec_idx: 0,
}
}
}
#[pyclass]
struct ScanRecordsIter {
owner: Py<ScanRecords>,
bucket_keys: Vec<TableBucket>,
bucket_idx: usize,
rec_idx: usize,
}
#[pymethods]
impl ScanRecordsIter {
fn __iter__(slf: PyRef<'_, Self>) -> PyRef<'_, Self> {
slf
}
fn __next__(&mut self, py: Python) -> Option<Py<ScanRecord>> {
let owner = self.owner.borrow(py);
loop {
if self.bucket_idx >= self.bucket_keys.len() {
return None;
}
let bucket = &self.bucket_keys[self.bucket_idx];
if let Some(recs) = owner.records_by_bucket.get(bucket) {
if self.rec_idx < recs.len() {
let rec = recs[self.rec_idx].clone_ref(py);
self.rec_idx += 1;
return Some(rec);
}
}
self.bucket_idx += 1;
self.rec_idx = 0;
}
}
}
/// Lazy iterator for `ScanRecords.items()` and `ScanRecords.values()`.
///
/// Yields one bucket at a time: `(TableBucket, list[ScanRecord])` for items,
/// or `list[ScanRecord]` for values. Only materializes records for the
/// current bucket on each `__next__` call.
#[pyclass]
pub struct ScanRecordsBucketIter {
owner: Py<ScanRecords>,
bucket_keys: Vec<TableBucket>,
bucket_idx: usize,
with_keys: bool,
}
#[pymethods]
impl ScanRecordsBucketIter {
fn __iter__(slf: PyRef<'_, Self>) -> PyRef<'_, Self> {
slf
}
fn __next__(&mut self, py: Python) -> Option<Py<PyAny>> {
if self.bucket_idx >= self.bucket_keys.len() {
return None;
}
let bucket = &self.bucket_keys[self.bucket_idx];
let owner = self.owner.borrow(py);
let recs = owner
.records_by_bucket
.get(bucket)
.map(|recs| recs.iter().map(|r| r.clone_ref(py)).collect::<Vec<_>>())
.unwrap_or_default();
let bucket = bucket.clone();
self.bucket_idx += 1;
if self.with_keys {
Some(
(bucket, recs)
.into_pyobject(py)
.unwrap()
.into_any()
.unbind(),
)
} else {
Some(recs.into_pyobject(py).unwrap().into_any().unbind())
}
}
}
/// Represents a Fluss table for data operations
#[pyclass]
pub struct FlussTable {
connection: Arc<fcore::client::FlussConnection>,
metadata: Arc<fcore::client::Metadata>,
table_info: fcore::metadata::TableInfo,
table_path: fcore::metadata::TablePath,
has_primary_key: bool,
}
/// Builder for creating log scanners with flexible configuration.
///
/// Use this builder to configure projection, and in the future, filters
/// before creating a log scanner.
#[pyclass]
pub struct TableScan {
connection: Arc<fcore::client::FlussConnection>,
metadata: Arc<fcore::client::Metadata>,
table_info: fcore::metadata::TableInfo,
projection: Option<ProjectionType>,
limit: Option<i32>,
}
/// Scanner type for internal use
enum ScannerType {
Record,
Batch,
}
#[pymethods]
impl TableScan {
/// Project to specific columns by their indices.
///
/// Args:
/// indices: List of column indices (0-based) to include in the scan.
///
/// Returns:
/// Self for method chaining.
pub fn project(mut slf: PyRefMut<'_, Self>, indices: Vec<usize>) -> PyRefMut<'_, Self> {
slf.projection = Some(ProjectionType::Indices(indices));
slf
}
/// Project to specific columns by their names.
///
/// Args:
/// names: List of column names to include in the scan.
///
/// Returns:
/// Self for method chaining.
pub fn project_by_name(mut slf: PyRefMut<'_, Self>, names: Vec<String>) -> PyRefMut<'_, Self> {
slf.projection = Some(ProjectionType::Names(names));
slf
}
/// Set a positive row limit, enabling `create_bucket_batch_scanner()`.
///
/// Args:
/// n: Maximum number of rows to scan. Must be positive.
///
/// Returns:
/// Self for method chaining.
pub fn limit(mut slf: PyRefMut<'_, Self>, n: i32) -> PyResult<PyRefMut<'_, Self>> {
if n <= 0 {
return Err(FlussError::new_err(format!(
"Scan limit must be positive, got {n}"
)));
}
slf.limit = Some(n);
Ok(slf)
}
/// Create a one-shot bounded scanner over a single bucket.
///
/// Requires a limit set via `limit()`; the scan runs on the first
/// `next_batch()`.
///
/// Args:
/// bucket: Bucket to scan; must belong to this table.
///
/// Returns:
/// A BatchScanner for `bucket`.
pub fn create_bucket_batch_scanner(&self, bucket: &TableBucket) -> PyResult<BatchScanner> {
let limit = self.limit.ok_or_else(|| {
FlussError::new_err("create_bucket_batch_scanner requires a limit set via .limit(n)")
})?;
let conn = self.connection.clone();
let _guard = TOKIO_RUNTIME.enter();
let table =
fcore::client::FlussTable::new(&conn, self.metadata.clone(), self.table_info.clone());
let projection = self.projection.clone();
let projection_indices = resolve_projection_indices(&projection, &self.table_info)?;
let scan = apply_projection(table.new_scan(), projection)?
.limit(limit)
.map_err(|e| FlussError::from_core_error(&e))?;
let scanner = scan
.create_bucket_batch_scanner(bucket.to_core())
.map_err(|e| FlussError::from_core_error(&e))?;
let (projected_schema, _) =
calculate_projected_types(&self.table_info, projection_indices)?;
Ok(BatchScanner::new(scanner, bucket.clone(), projected_schema))
}
/// Create a record-based log scanner.
///
/// Use this scanner with `poll()` to get individual records with metadata
/// (offset, timestamp, change_type).
///
/// Returns:
/// LogScanner for record-by-record scanning with `poll()`
pub fn create_log_scanner<'py>(&self, py: Python<'py>) -> PyResult<Bound<'py, PyAny>> {
self.create_scanner_internal(py, ScannerType::Record)
}
/// Create a batch-based log scanner.
///
/// Use this scanner with `poll_arrow()` to get Arrow Tables, or with
/// `poll_record_batch()` to get individual batches with metadata.
///
/// Returns:
/// LogScanner for batch-based scanning with `poll_arrow()` or `poll_record_batch()`
pub fn create_record_batch_log_scanner<'py>(
&self,
py: Python<'py>,
) -> PyResult<Bound<'py, PyAny>> {
self.create_scanner_internal(py, ScannerType::Batch)
}
fn __repr__(&self) -> String {
format!(
"TableScan(table={}.{})",
self.table_info.table_path.database(),
self.table_info.table_path.table()
)
}
}
impl TableScan {
fn create_scanner_internal<'py>(
&self,
py: Python<'py>,
scanner_type: ScannerType,
) -> PyResult<Bound<'py, PyAny>> {
if let Some(limit) = self.limit {
return Err(FlussError::new_err(format!(
"Log scanners don't support limit pushdown (requested limit: {limit}). \
Use create_bucket_batch_scanner() for a bounded scan."
)));
}
let conn = self.connection.clone();
let metadata = self.metadata.clone();
let table_info = self.table_info.clone();
let projection = self.projection.clone();
future_into_py(py, async move {
let fluss_table = fcore::client::FlussTable::new(&conn, metadata, table_info.clone());
let projection_indices = resolve_projection_indices(&projection, &table_info)?;
let table_scan = apply_projection(fluss_table.new_scan(), projection)?;
let admin = conn
.get_admin()
.map_err(|e| FlussError::from_core_error(&e))?;
let (projected_schema, projected_row_type) =
calculate_projected_types(&table_info, projection_indices)?;
let scanner_kind = match scanner_type {
ScannerType::Record => {
let s = table_scan
.create_log_scanner()
.map_err(|e| FlussError::from_core_error(&e))?;
ScannerKind::Record(s)
}
ScannerType::Batch => {
let s = table_scan
.create_record_batch_log_scanner()
.map_err(|e| FlussError::from_core_error(&e))?;
ScannerKind::Batch(s)
}
};
let py_scanner = LogScanner::new(
scanner_kind,
admin,
table_info,
projected_schema,
Arc::new(projected_row_type),
);
Python::attach(|py| Py::new(py, py_scanner))
})
}
}
/// Internal enum to represent different projection types
#[derive(Clone)]
enum ProjectionType {
Indices(Vec<usize>),
Names(Vec<String>),
}
/// Resolve projection to column indices
fn resolve_projection_indices(
projection: &Option<ProjectionType>,
table_info: &fcore::metadata::TableInfo,
) -> PyResult<Option<Vec<usize>>> {
match projection {
Some(ProjectionType::Indices(indices)) => Ok(Some(indices.clone())),
Some(ProjectionType::Names(names)) => {
let schema = table_info.get_schema();
let columns = schema.columns();
let mut indices = Vec::with_capacity(names.len());
for name in names {
let idx = columns
.iter()
.position(|c| c.name() == name)
.ok_or_else(|| FlussError::new_err(format!("Column '{name}' not found")))?;
indices.push(idx);
}
Ok(Some(indices))
}
None => Ok(None),
}
}
/// Apply projection to table scan
fn apply_projection(
table_scan: fcore::client::TableScan,
projection: Option<ProjectionType>,
) -> PyResult<fcore::client::TableScan> {
match projection {
Some(ProjectionType::Indices(indices)) => table_scan
.project(&indices)
.map_err(|e| FlussError::from_core_error(&e)),
Some(ProjectionType::Names(names)) => {
let column_name_refs: Vec<&str> = names.iter().map(|s| s.as_str()).collect();
table_scan
.project_by_name(&column_name_refs)
.map_err(|e| FlussError::from_core_error(&e))
}
None => Ok(table_scan),
}
}
/// Calculate projected schema and row type from projection indices
fn calculate_projected_types(
table_info: &fcore::metadata::TableInfo,
projection_indices: Option<Vec<usize>>,
) -> PyResult<(SchemaRef, fcore::metadata::RowType)> {
let full_schema =
to_arrow_schema(table_info.get_row_type()).map_err(|e| FlussError::from_core_error(&e))?;
let full_row_type = table_info.get_row_type();
match projection_indices {
Some(indices) => {
let arrow_fields: Vec<_> = indices
.iter()
.map(|&i| full_schema.field(i).clone())
.collect();
let row_fields: Vec<_> = indices
.iter()
.map(|&i| full_row_type.fields()[i].clone())
.collect();
Ok((
Arc::new(arrow_schema::Schema::new(arrow_fields)),
fcore::metadata::RowType::new(row_fields),
))
}
None => Ok((full_schema, full_row_type.clone())),
}
}
#[pymethods]
impl FlussTable {
/// Create a new table scan builder for configuring and creating log scanners.
///
/// Use this method to create scanners with the builder pattern:
/// Returns:
/// TableScan builder for configuring the scanner.
pub fn new_scan(&self) -> TableScan {
TableScan {
connection: self.connection.clone(),
metadata: self.metadata.clone(),
table_info: self.table_info.clone(),
projection: None,
limit: None,
}
}
/// Create a new TableAppend builder for the table.
///
/// Returns:
/// TableAppend builder. Call `create_writer()` to get an AppendWriter.
fn new_append(&self) -> PyResult<TableAppend> {
let _guard = TOKIO_RUNTIME.enter();
let fluss_table = fcore::client::FlussTable::new(
&self.connection,
self.metadata.clone(),
self.table_info.clone(),
);
let table_append = fluss_table
.new_append()
.map_err(|e| FlussError::from_core_error(&e))?;
Ok(TableAppend {
inner: table_append,
table_info: self.table_info.clone(),
})
}
/// Get table information
pub fn get_table_info(&self) -> TableInfo {
TableInfo::from_core(self.table_info.clone())
}
/// Get table path
pub fn get_table_path(&self) -> TablePath {
TablePath::from_core(self.table_path.clone())
}
/// Check if table has primary key
pub fn has_primary_key(&self) -> bool {
self.has_primary_key
}
/// Create a new TableLookup builder for primary key lookups.
///
/// This is only available for tables with a primary key.
///
/// Returns:
/// TableLookup builder. Call `create_lookuper()` to get a Lookuper.
pub fn new_lookup(&self) -> PyResult<TableLookup> {
if !self.has_primary_key {
return Err(FlussError::new_err(
"Lookup is only supported for primary key tables",
));
}
Ok(TableLookup {
connection: self.connection.clone(),
metadata: self.metadata.clone(),
table_info: self.table_info.clone(),
})
}
/// Create a new TableUpsert builder for the table.
///
/// This is only available for tables with a primary key.
///
/// Returns:
/// TableUpsert builder. Call `create_writer()` to get an UpsertWriter,
/// or use `partial_update_by_name()` / `partial_update_by_index()` first.
pub fn new_upsert(&self) -> PyResult<TableUpsert> {
if !self.has_primary_key {
return Err(FlussError::new_err(
"Upsert is only supported for primary key tables",
));
}
let _guard = TOKIO_RUNTIME.enter();
let fluss_table = fcore::client::FlussTable::new(
&self.connection,
self.metadata.clone(),
self.table_info.clone(),
);
let table_upsert = fluss_table
.new_upsert()
.map_err(|e| FlussError::from_core_error(&e))?;
Ok(TableUpsert {
inner: table_upsert,
table_info: self.table_info.clone(),
target_columns: None,
})
}
fn __repr__(&self) -> String {
format!(
"FlussTable(path={}.{})",
self.table_path.database(),
self.table_path.table()
)
}
}
impl FlussTable {
/// Create a FlussTable
pub fn new_table(
connection: Arc<fcore::client::FlussConnection>,
metadata: Arc<fcore::client::Metadata>,
table_info: fcore::metadata::TableInfo,
table_path: fcore::metadata::TablePath,
has_primary_key: bool,
) -> Self {
Self {
connection,
metadata,
table_info,
table_path,
has_primary_key,
}
}
}
/// Builder for creating an AppendWriter.
///
/// Obtain via `FlussTable.new_append()`, then call `create_writer()`.
#[pyclass]
pub struct TableAppend {
inner: fcore::client::TableAppend,
table_info: fcore::metadata::TableInfo,
}
#[pymethods]
impl TableAppend {
/// Create an AppendWriter from this builder.
pub fn create_writer(&self) -> PyResult<AppendWriter> {
let rust_writer = self
.inner
.create_writer()
.map_err(|e| FlussError::from_core_error(&e))?;
Ok(AppendWriter::from_core(
rust_writer,
self.table_info.clone(),
))
}
fn __repr__(&self) -> String {
"TableAppend()".to_string()
}
}
/// Builder for creating an UpsertWriter, with optional partial update configuration.
///
/// Obtain via `FlussTable.new_upsert()`, then optionally call
/// `partial_update_by_name()` or `partial_update_by_index()`,
/// then call `create_writer()`.
#[pyclass]
pub struct TableUpsert {
inner: fcore::client::TableUpsert,
table_info: fcore::metadata::TableInfo,
/// Column indices for partial updates, tracked for Python's dict→GenericRow conversion.
target_columns: Option<Vec<usize>>,
}
#[pymethods]
impl TableUpsert {
/// Configure partial update by column names.
///
/// Only the specified columns will be updated on upsert.
///
/// Args:
/// columns: List of column names to update.
///
/// Returns:
/// A new TableUpsert configured for partial update.
pub fn partial_update_by_name(&self, columns: Vec<String>) -> PyResult<TableUpsert> {
let col_refs: Vec<&str> = columns.iter().map(|s| s.as_str()).collect();
// Core validates and resolves names → indices internally
let updated = self
.inner
.partial_update_with_column_names(&col_refs)
.map_err(|e| FlussError::from_core_error(&e))?;
// Resolve indices for Python's row conversion layer (core validated names above)
let row_type = self.table_info.row_type();
let indices: Vec<usize> = columns
.iter()
.map(|name| {
row_type.get_field_index(name).ok_or_else(|| {
FlussError::new_err(format!("Unknown column name '{name}' for partial update"))
})
})
.collect::<PyResult<Vec<usize>>>()?;
Ok(TableUpsert {
inner: updated,
table_info: self.table_info.clone(),
target_columns: Some(indices),
})
}
/// Configure partial update by column indices.
///
/// Only the specified columns will be updated on upsert.
///
/// Args:
/// column_indices: List of column indices (0-based) to update.
///
/// Returns:
/// A new TableUpsert configured for partial update.
pub fn partial_update_by_index(&self, column_indices: Vec<usize>) -> PyResult<TableUpsert> {
let target = column_indices.clone();
// Core validates indices internally
let updated = self
.inner
.partial_update(Some(column_indices))
.map_err(|e| FlussError::from_core_error(&e))?;
Ok(TableUpsert {
inner: updated,
table_info: self.table_info.clone(),
target_columns: Some(target),
})
}
/// Create an UpsertWriter from this builder.
pub fn create_writer(&self) -> PyResult<crate::UpsertWriter> {
crate::UpsertWriter::new(
&self.inner,
self.table_info.clone(),
self.target_columns.clone(),
)
}
fn __repr__(&self) -> String {
"TableUpsert()".to_string()
}
}
/// Builder for creating a Lookuper.
///
/// Obtain via `FlussTable.new_lookup()`, then call `create_lookuper()`.
#[pyclass]
pub struct TableLookup {
connection: Arc<fcore::client::FlussConnection>,
metadata: Arc<fcore::client::Metadata>,
table_info: fcore::metadata::TableInfo,
}
#[pymethods]
impl TableLookup {
/// Create a Lookuper from this builder.
pub fn create_lookuper(&self) -> PyResult<crate::Lookuper> {
crate::Lookuper::new(
&self.connection,
self.metadata.clone(),
self.table_info.clone(),
)
}
/// Switch to prefix-scan mode for the given lookup columns.
///
/// The columns must be the table's partition keys (if any) plus the
/// bucket keys, in that order.
///
/// Args:
/// column_names: List of column names forming the prefix key.
///
/// Returns:
/// TablePrefixLookup builder. Call `create_lookuper()` to get a PrefixLookuper.
pub fn lookup_by(&self, column_names: Vec<String>) -> TablePrefixLookup {
TablePrefixLookup {
connection: self.connection.clone(),
metadata: self.metadata.clone(),
table_info: self.table_info.clone(),
lookup_column_names: column_names,
}
}
fn __repr__(&self) -> String {
"TableLookup()".to_string()
}
}
/// Builder for creating a PrefixLookuper.
///
/// Obtain via `TableLookup.lookup_by(columns)`, then call `create_lookuper()`.
#[pyclass]
pub struct TablePrefixLookup {
connection: Arc<fcore::client::FlussConnection>,
metadata: Arc<fcore::client::Metadata>,
table_info: fcore::metadata::TableInfo,
lookup_column_names: Vec<String>,
}
#[pymethods]
impl TablePrefixLookup {
/// Create a PrefixLookuper from this builder.
pub fn create_lookuper(&self) -> PyResult<crate::PrefixLookuper> {
crate::PrefixLookuper::new(
&self.connection,
self.metadata.clone(),
self.table_info.clone(),
self.lookup_column_names.clone(),
)
}
fn __repr__(&self) -> String {
"TablePrefixLookup()".to_string()
}
}
/// Writer for appending data to a Fluss table
#[pyclass]
pub struct AppendWriter {
inner: Arc<fcore::client::AppendWriter>,
table_info: fcore::metadata::TableInfo,
}
#[pymethods]
impl AppendWriter {
/// Write Arrow table data (fire-and-forget, use flush() to ensure delivery)
pub fn write_arrow(&self, py: Python, table: Py<PyAny>) -> PyResult<()> {
// Convert Arrow Table to batches and write each batch
let batches = table.call_method0(py, "to_batches")?;
let batch_list: Vec<Py<PyAny>> = batches.extract(py)?;
for batch in batch_list {
// Drop the handle — fire-and-forget for bulk writes
drop(self.write_arrow_batch(py, batch)?);
}
Ok(())
}
/// Write Arrow batch data.
///
/// Returns:
/// WriteResultHandle that can be ignored (fire-and-forget) or
/// awaited via `handle.wait()` for server acknowledgment.
pub fn write_arrow_batch(&self, py: Python, batch: Py<PyAny>) -> PyResult<WriteResultHandle> {
// This shares the underlying Arrow buffers without copying data
let batch_bound = batch.bind(py);
let rust_batch: ArrowRecordBatch = FromPyArrow::from_pyarrow_bound(batch_bound)
.map_err(|e| FlussError::new_err(format!("Failed to convert RecordBatch: {e}")))?;
let result_future = self
.inner
.append_arrow_batch(rust_batch)
.map_err(|e| FlussError::from_core_error(&e))?;
Ok(WriteResultHandle::new(result_future))
}
/// Append a single row to the table.
///
/// Returns:
/// WriteResultHandle that can be ignored (fire-and-forget) or
/// awaited via `handle.wait()` for server acknowledgment.
pub fn append(&self, row: &Bound<'_, PyAny>) -> PyResult<WriteResultHandle> {
let generic_row = python_to_generic_row(row, &self.table_info)?;
let result_future = self
.inner
.append(&generic_row)
.map_err(|e| FlussError::from_core_error(&e))?;
Ok(WriteResultHandle::new(result_future))
}
/// Write Pandas DataFrame data
pub fn write_pandas(&self, py: Python, df: Py<PyAny>) -> PyResult<()> {
// Get the expected Arrow schema from the Fluss table
let row_type = self.table_info.get_row_type();
let expected_schema = fcore::record::to_arrow_schema(row_type)
.map_err(|e| FlussError::from_core_error(&e))?;
// Convert Arrow schema to PyArrow schema
let py_schema = expected_schema
.as_ref()
.to_pyarrow(py)
.map_err(|e| FlussError::new_err(format!("Failed to convert schema: {e}")))?;
// Import pyarrow module
let pyarrow = py.import("pyarrow")?;
// Get the Table class from pyarrow module
let table_class = pyarrow.getattr("Table")?;
// Call Table.from_pandas(df, schema=expected_schema) to ensure proper type casting
let pa_table = table_class.call_method(
"from_pandas",
(df,),
Some(&[("schema", py_schema)].into_py_dict(py)?),
)?;
// Then call write_arrow with the converted table
self.write_arrow(py, pa_table.into())
}
/// Flush any pending data
pub fn flush<'py>(&self, py: Python<'py>) -> PyResult<Bound<'py, PyAny>> {
let inner = self.inner.clone();
future_into_py(py, async move {
inner
.flush()
.await
.map_err(|e| FlussError::from_core_error(&e))
})
}
// Enter the async runtime context (for 'async with' statement)
fn __aenter__<'py>(slf: PyRef<'py, Self>, py: Python<'py>) -> PyResult<Bound<'py, PyAny>> {
let py_slf = slf.into_pyobject(py)?.unbind();
future_into_py(py, async move { Ok(py_slf) })
}
// Exit the async runtime context (for 'async with' statement)
/// On exit, the writer is automatically flushed.
#[pyo3(signature = (exc_type=None, _exc_value=None, _traceback=None))]
fn __aexit__<'py>(
&self,
py: Python<'py>,
exc_type: Option<Bound<'py, PyAny>>,
_exc_value: Option<Bound<'py, PyAny>>,
_traceback: Option<Bound<'py, PyAny>>,
) -> PyResult<Bound<'py, PyAny>> {
let inner = self.inner.clone();
let is_exc_none = exc_type.as_ref().is_none_or(|e| e.is_none());
future_into_py(py, async move {
let res = inner.flush().await;
if let Err(e) = res {
if is_exc_none {
return Err(FlussError::from_core_error(&e));
}
}
Ok(false)
})
}
fn __repr__(&self) -> String {
"AppendWriter()".to_string()
}
}
impl AppendWriter {
/// Create a AppendWriter from a core append writer
pub fn from_core(
append: fcore::client::AppendWriter,
table_info: fcore::metadata::TableInfo,
) -> Self {
Self {
inner: Arc::new(append),
table_info,
}
}
}
/// Represents different input shapes for a row
#[derive(FromPyObject)]
enum RowInput<'py> {
Dict(Bound<'py, PyDict>),
Tuple(Bound<'py, PyTuple>),
List(Bound<'py, PyList>),
}
/// Convert Python row (dict/list/tuple) to GenericRow requiring all schema columns.
pub fn python_to_generic_row(
row: &Bound<PyAny>,
table_info: &fcore::metadata::TableInfo,
) -> PyResult<fcore::row::GenericRow<'static>> {
let all_indices: Vec<usize> = (0..table_info.row_type().fields().len()).collect();
python_to_sparse_generic_row(row, table_info, &all_indices)
}
/// Process a Python sequence (list or tuple) into datums at the target column positions.
fn process_sequence(
seq: &Bound<PySequence>,
target_indices: &[usize],
fields: &[fcore::metadata::DataField],
datums: &mut [fcore::row::Datum<'static>],
sparse: bool,
) -> PyResult<()> {
if seq.len()? != target_indices.len() {
return Err(FlussError::new_err(format!(
"Expected {} elements, got {}",
target_indices.len(),
seq.len()?
)));
}
for (i, &col_idx) in target_indices.iter().enumerate() {
let field = &fields[col_idx];
let value = seq.get_item(i)?;
let dest = if sparse { col_idx } else { i };
datums[dest] = python_value_to_datum(&value, field.data_type())
.map_err(|e| FlussError::new_err(format!("Field '{}': {}", field.name(), e)))?;
}
Ok(())
}
/// Build a full-width GenericRow filling only the specified column
/// indices from user input; all other columns are set to Null.
pub fn python_to_sparse_generic_row(
row: &Bound<PyAny>,
table_info: &fcore::metadata::TableInfo,
target_indices: &[usize],
) -> PyResult<fcore::row::GenericRow<'static>> {
python_to_generic_row_inner(row, table_info, target_indices, true)
}
/// Build a dense GenericRow with exactly `target_indices.len()` fields,
/// containing only the target column values in order.
pub fn python_to_dense_generic_row(
row: &Bound<PyAny>,
table_info: &fcore::metadata::TableInfo,
target_indices: &[usize],
) -> PyResult<fcore::row::GenericRow<'static>> {
python_to_generic_row_inner(row, table_info, target_indices, false)
}
/// Build a GenericRow from user input. When `sparse` is true, the row is full width and padded with nulls
fn python_to_generic_row_inner(
row: &Bound<PyAny>,
table_info: &fcore::metadata::TableInfo,
target_indices: &[usize],
sparse: bool,
) -> PyResult<fcore::row::GenericRow<'static>> {
let row_type = table_info.row_type();
let fields = row_type.fields();
let target_names: Vec<&str> = target_indices.iter().map(|&i| fields[i].name()).collect();
let num_fields = if sparse {
fields.len()
} else {
target_indices.len()
};
let mut datums: Vec<fcore::row::Datum<'static>> = vec![fcore::row::Datum::Null; num_fields];
let row_input: RowInput = row.extract().map_err(|_| {
let type_name = row
.get_type()
.name()
.map(|n| n.to_string())
.unwrap_or_else(|_| "unknown".to_string());
FlussError::new_err(format!(
"Row must be a dict, list, or tuple; got {type_name}"
))
})?;
match row_input {
RowInput::Dict(dict) => {
for (k, _) in dict.iter() {
let key_str = k.extract::<&str>().map_err(|_| {
let key_type = k
.get_type()
.name()
.map(|n| n.to_string())
.unwrap_or_else(|_| "unknown".to_string());
FlussError::new_err(format!("Dict keys must be strings; got {key_type}"))
})?;
if !target_names.contains(&key_str) {
return Err(FlussError::new_err(format!(
"Unknown field '{}'. Expected: {}",
key_str,
target_names.join(", ")
)));
}
}
for (i, &col_idx) in target_indices.iter().enumerate() {
let name = target_names[i];
let field = &fields[col_idx];
let dest = if sparse { col_idx } else { i };
match dict.get_item(name)? {
Some(value) => {
datums[dest] = python_value_to_datum(&value, field.data_type())
.map_err(|e| FlussError::new_err(format!("Field '{name}': {e}")))?;
}
None if field.data_type().is_nullable() => {}
None => {
return Err(FlussError::new_err(format!(
"Missing value for non-nullable field '{name}'"
)));
}
}
}
}
RowInput::List(list) => {
process_sequence(
list.as_sequence(),
target_indices,
fields,
&mut datums,
sparse,
)?;
}
RowInput::Tuple(tuple) => {
process_sequence(
tuple.as_sequence(),
target_indices,
fields,
&mut datums,
sparse,
)?;
}
}
Ok(fcore::row::GenericRow { values: datums })
}
/// Convert Python value to Datum based on data type
fn python_value_to_datum(value: &Bound<PyAny>, data_type: &DataType) -> PyResult<Datum<'static>> {
if value.is_none() {
return Ok(Datum::Null);
}
match data_type {
DataType::Boolean(_) => {
let v: bool = value.extract()?;
Ok(Datum::Bool(v))
}
DataType::TinyInt(_) => {
// Strict type checking: reject bool for int columns
if value.is_instance_of::<PyBool>() {
return Err(FlussError::new_err(
"Expected int for TinyInt column, got bool. Use 0 or 1 explicitly.".to_string(),
));
}
let v: i8 = value.extract()?;
Ok(Datum::Int8(v))
}
DataType::SmallInt(_) => {
if value.is_instance_of::<PyBool>() {
return Err(FlussError::new_err(
"Expected int for SmallInt column, got bool. Use 0 or 1 explicitly."
.to_string(),
));
}
let v: i16 = value.extract()?;
Ok(Datum::Int16(v))
}
DataType::Int(_) => {
if value.is_instance_of::<PyBool>() {
return Err(FlussError::new_err(
"Expected int for Int column, got bool. Use 0 or 1 explicitly.".to_string(),
));
}
let v: i32 = value.extract()?;
Ok(Datum::Int32(v))
}
DataType::BigInt(_) => {
if value.is_instance_of::<PyBool>() {
return Err(FlussError::new_err(
"Expected int for BigInt column, got bool. Use 0 or 1 explicitly.".to_string(),
));
}
let v: i64 = value.extract()?;
Ok(Datum::Int64(v))
}
DataType::Float(_) => {
let v: f32 = value.extract()?;
Ok(Datum::Float32(F32::from(v)))
}
DataType::Double(_) => {
let v: f64 = value.extract()?;
Ok(Datum::Float64(F64::from(v)))
}
DataType::String(_) | DataType::Char(_) => {
let v: String = value.extract()?;
Ok(v.into())
}
DataType::Bytes(_) | DataType::Binary(_) => {
// Efficient extraction: downcast to specific type and use bulk copy.
// PyBytes::as_bytes() and PyByteArray::to_vec() are O(n) bulk copies of the underlying data.
if let Ok(bytes) = value.downcast::<PyBytes>() {
Ok(bytes.as_bytes().to_vec().into())
} else if let Ok(bytearray) = value.downcast::<PyByteArray>() {
Ok(bytearray.to_vec().into())
} else {
Err(FlussError::new_err(format!(
"Expected bytes or bytearray, got {}",
value.get_type().name()?
)))
}
}
DataType::Decimal(decimal_type) => {
python_decimal_to_datum(value, decimal_type.precision(), decimal_type.scale())
}
DataType::Date(_) => python_date_to_datum(value),
DataType::Time(_) => python_time_to_datum(value),
DataType::Timestamp(_) => python_datetime_to_timestamp_ntz(value),
DataType::TimestampLTz(_) => python_datetime_to_timestamp_ltz(value),
DataType::Array(array_type) => {
let element_type = array_type.get_element_type();
if value.is_instance_of::<PyString>() {
return Err(FlussError::new_err(format!(
"Expected sequence for Array column, got {}",
get_type_name(value)
)));
}
let seq = value.downcast::<PySequence>().map_err(|_| {
FlussError::new_err(format!(
"Expected sequence for Array column, got {}",
get_type_name(value)
))
})?;
let len = seq.len()?;
let mut writer = FlussArrayWriter::new(len, element_type);
for i in 0..len {
let item = seq.get_item(i)?;
if item.is_none() {
writer.set_null_at(i);
} else {
let val_datum = python_value_to_datum(&item, element_type)?;
match val_datum {
Datum::Null => writer.set_null_at(i),
Datum::Bool(v) => writer.write_boolean(i, v),
Datum::Int8(v) => writer.write_byte(i, v),
Datum::Int16(v) => writer.write_short(i, v),
Datum::Int32(v) => writer.write_int(i, v),
Datum::Int64(v) => writer.write_long(i, v),
Datum::Float32(v) => writer.write_float(i, v.into_inner()),
Datum::Float64(v) => writer.write_double(i, v.into_inner()),
Datum::String(v) => writer.write_string(i, &v),
Datum::Blob(v) => writer.write_binary_bytes(i, v.as_ref()),
Datum::Decimal(v) => {
if let DataType::Decimal(dt) = element_type {
writer.write_decimal(i, &v, dt.precision());
}
}
Datum::Date(v) => writer.write_date(i, v),
Datum::Time(v) => writer.write_time(i, v),
Datum::TimestampNtz(v) => {
if let DataType::Timestamp(dt) = element_type {
writer.write_timestamp_ntz(i, &v, dt.precision());
}
}
Datum::TimestampLtz(v) => {
if let DataType::TimestampLTz(dt) = element_type {
writer.write_timestamp_ltz(i, &v, dt.precision());
}
}
Datum::Array(v) => writer.write_array(i, &v),
Datum::Map(v) => writer.write_map(i, &v),
Datum::Row(v) => writer
.write_row(i, v.as_ref())
.map_err(|e| FlussError::from_core_error(&e))?,
}
}
}
let array = writer
.complete()
.map_err(|e| FlussError::from_core_error(&e))?;
Ok(Datum::Array(array))
}
DataType::Map(map_type) => {
let key_type = map_type.key_type();
let value_type = map_type.value_type();
let pairs = python_map_pairs(value)?;
let mut writer = FlussMapWriter::new(pairs.len(), key_type, value_type);
for (k, v) in pairs {
let key_datum = python_value_to_datum(&k, key_type)?;
let value_datum = python_value_to_datum(&v, value_type)?;
writer
.write_entry(key_datum, value_datum)
.map_err(|e| FlussError::from_core_error(&e))?;
}
let map = writer
.complete()
.map_err(|e| FlussError::from_core_error(&e))?;
Ok(Datum::Map(map))
}
DataType::Row(row_type) => {
let nested = python_value_to_row_datum(value, row_type)?;
Ok(Datum::Row(Box::new(nested)))
}
}
}
/// Extract `(key, value)` pairs from a Python value representing a MAP column.
/// Accepts a `dict`, or a sequence of `(key, value)` pairs — matching the
/// shape pyarrow's `MapArray.to_pylist()` produces, so MAP values round-trip.
fn python_map_pairs<'py>(
value: &Bound<'py, PyAny>,
) -> PyResult<Vec<(Bound<'py, PyAny>, Bound<'py, PyAny>)>> {
if let Ok(dict) = value.downcast::<PyDict>() {
return Ok(dict.iter().collect());
}
if value.is_instance_of::<PyString>() {
return Err(FlussError::new_err(format!(
"Expected dict or sequence of (key, value) pairs for Map column, got {}",
get_type_name(value)
)));
}
let seq = value.downcast::<PySequence>().map_err(|_| {
FlussError::new_err(format!(
"Expected dict or sequence of (key, value) pairs for Map column, got {}",
get_type_name(value)
))
})?;
let len = seq.len()?;
let mut pairs = Vec::with_capacity(len);
for i in 0..len {
let entry = seq.get_item(i)?;
let pair = entry.downcast::<PySequence>().map_err(|_| {
FlussError::new_err("Map entries must be (key, value) pairs".to_string())
})?;
if pair.len()? != 2 {
return Err(FlussError::new_err(
"Map entries must be (key, value) pairs of length 2".to_string(),
));
}
pairs.push((pair.get_item(0)?, pair.get_item(1)?));
}
Ok(pairs)
}
/// Convert a Python value (`dict` by field name, or `list`/`tuple` by position)
/// into a nested `GenericRow` for a ROW column.
fn python_value_to_row_datum(
value: &Bound<PyAny>,
row_type: &RowType,
) -> PyResult<GenericRow<'static>> {
let fields = row_type.fields();
let mut datums: Vec<Datum<'static>> = vec![Datum::Null; fields.len()];
let row_input: RowInput = value.extract().map_err(|_| {
FlussError::new_err(format!(
"Row column must be a dict, list, or tuple; got {}",
get_type_name(value)
))
})?;
match row_input {
RowInput::Dict(dict) => {
for (k, _) in dict.iter() {
let key_str = k.extract::<&str>().map_err(|_| {
FlussError::new_err(format!(
"Row field keys must be strings; got {}",
get_type_name(&k)
))
})?;
if !fields.iter().any(|f| f.name() == key_str) {
return Err(FlussError::new_err(format!(
"Unknown row field '{}'. Expected: {}",
key_str,
fields
.iter()
.map(|f| f.name())
.collect::<Vec<_>>()
.join(", ")
)));
}
}
for (i, field) in fields.iter().enumerate() {
match dict.get_item(field.name())? {
Some(v) => {
datums[i] = python_value_to_datum(&v, field.data_type()).map_err(|e| {
FlussError::new_err(format!("Row field '{}': {}", field.name(), e))
})?;
}
None if field.data_type().is_nullable() => {}
None => {
return Err(FlussError::new_err(format!(
"Missing value for non-nullable row field '{}'",
field.name()
)));
}
}
}
}
RowInput::List(list) => fill_row_fields(list.as_sequence(), fields, &mut datums)?,
RowInput::Tuple(tuple) => fill_row_fields(tuple.as_sequence(), fields, &mut datums)?,
}
Ok(GenericRow { values: datums })
}
/// Fill ROW field datums from a positional Python sequence.
fn fill_row_fields(
seq: &Bound<PySequence>,
fields: &[DataField],
datums: &mut [Datum<'static>],
) -> PyResult<()> {
if seq.len()? != fields.len() {
return Err(FlussError::new_err(format!(
"Expected {} row fields, got {}",
fields.len(),
seq.len()?
)));
}
for (i, field) in fields.iter().enumerate() {
let v = seq.get_item(i)?;
datums[i] = python_value_to_datum(&v, field.data_type())
.map_err(|e| FlussError::new_err(format!("Row field '{}': {}", field.name(), e)))?;
}
Ok(())
}
/// Convert Rust Datum to Python value based on data type.
/// This is the reverse of python_value_to_datum.
pub fn datum_to_python_value(
py: Python,
row: &dyn InternalRow,
pos: usize,
data_type: &DataType,
) -> PyResult<Py<PyAny>> {
// Check for null first
if row
.is_null_at(pos)
.map_err(|e| FlussError::from_core_error(&e))?
{
return Ok(py.None());
}
match data_type {
DataType::Boolean(_) => Ok(row
.get_boolean(pos)
.map_err(|e| FlussError::from_core_error(&e))?
.into_pyobject(py)?
.to_owned()
.into_any()
.unbind()),
DataType::TinyInt(_) => Ok(row
.get_byte(pos)
.map_err(|e| FlussError::from_core_error(&e))?
.into_pyobject(py)?
.to_owned()
.into_any()
.unbind()),
DataType::SmallInt(_) => Ok(row
.get_short(pos)
.map_err(|e| FlussError::from_core_error(&e))?
.into_pyobject(py)?
.to_owned()
.into_any()
.unbind()),
DataType::Int(_) => Ok(row
.get_int(pos)
.map_err(|e| FlussError::from_core_error(&e))?
.into_pyobject(py)?
.to_owned()
.into_any()
.unbind()),
DataType::BigInt(_) => Ok(row
.get_long(pos)
.map_err(|e| FlussError::from_core_error(&e))?
.into_pyobject(py)?
.to_owned()
.into_any()
.unbind()),
DataType::Float(_) => Ok(row
.get_float(pos)
.map_err(|e| FlussError::from_core_error(&e))?
.into_pyobject(py)?
.to_owned()
.into_any()
.unbind()),
DataType::Double(_) => Ok(row
.get_double(pos)
.map_err(|e| FlussError::from_core_error(&e))?
.into_pyobject(py)?
.to_owned()
.into_any()
.unbind()),
DataType::String(_) => {
let s = row
.get_string(pos)
.map_err(|e| FlussError::from_core_error(&e))?;
Ok(s.into_pyobject(py)?.into_any().unbind())
}
DataType::Char(char_type) => {
let s = row
.get_char(pos, char_type.length() as usize)
.map_err(|e| FlussError::from_core_error(&e))?;
Ok(s.into_pyobject(py)?.into_any().unbind())
}
DataType::Bytes(_) => {
let b = row
.get_bytes(pos)
.map_err(|e| FlussError::from_core_error(&e))?;
Ok(PyBytes::new(py, b).into_any().unbind())
}
DataType::Binary(binary_type) => {
let b = row
.get_binary(pos, binary_type.length())
.map_err(|e| FlussError::from_core_error(&e))?;
Ok(PyBytes::new(py, b).into_any().unbind())
}
DataType::Decimal(decimal_type) => {
let decimal = row
.get_decimal(
pos,
decimal_type.precision() as usize,
decimal_type.scale() as usize,
)
.map_err(|e| FlussError::from_core_error(&e))?;
rust_decimal_to_python(py, &decimal)
}
DataType::Date(_) => {
let date = row
.get_date(pos)
.map_err(|e| FlussError::from_core_error(&e))?;
rust_date_to_python(py, date)
}
DataType::Time(_) => {
let time = row
.get_time(pos)
.map_err(|e| FlussError::from_core_error(&e))?;
rust_time_to_python(py, time)
}
DataType::Timestamp(ts_type) => {
let ts = row
.get_timestamp_ntz(pos, ts_type.precision())
.map_err(|e| FlussError::from_core_error(&e))?;
rust_timestamp_ntz_to_python(py, ts)
}
DataType::TimestampLTz(ts_type) => {
let ts = row
.get_timestamp_ltz(pos, ts_type.precision())
.map_err(|e| FlussError::from_core_error(&e))?;
rust_timestamp_ltz_to_python(py, ts)
}
DataType::Array(array_type) => {
let array_data = row
.get_array(pos)
.map_err(|e| FlussError::from_core_error(&e))?
.try_into_binary()
.map_err(|e| FlussError::from_core_error(&e))?;
array_to_pylist(py, &array_data, array_type.get_element_type())
}
DataType::Map(map_type) => {
let map_data = row
.get_map(pos)
.map_err(|e| FlussError::from_core_error(&e))?
.try_into_binary()
.map_err(|e| FlussError::from_core_error(&e))?;
map_to_pylist(py, &map_data, map_type)
}
DataType::Row(row_type) => {
let nested = row
.get_row(pos)
.map_err(|e| FlussError::from_core_error(&e))?
.try_into_generic(row_type)
.map_err(|e| FlussError::from_core_error(&e))?;
row_to_pydict(py, &nested, row_type)
}
}
}
/// Convert a binary `FlussArray` to a Python list.
fn array_to_pylist(py: Python, arr: &FlussArray, element_type: &DataType) -> PyResult<Py<PyAny>> {
let py_list = pyo3::types::PyList::empty(py);
for i in 0..arr.size() {
py_list.append(array_elem_to_python(py, arr, i, element_type)?)?;
}
Ok(py_list.into_any().unbind())
}
/// Convert a Fluss `MAP` to a Python list of `(key, value)` tuples — matching
/// pyarrow's default `MapArray.to_pylist()` shape (preserves duplicate keys and
/// ordering, and allows non-hashable keys).
fn map_to_pylist(py: Python, map_data: &FlussMap, map_type: &MapType) -> PyResult<Py<PyAny>> {
let keys = map_data.key_array();
let values = map_data.value_array();
let py_list = pyo3::types::PyList::empty(py);
for i in 0..map_data.size() {
let py_key = array_elem_to_python(py, keys, i, map_type.key_type())?;
let py_val = array_elem_to_python(py, values, i, map_type.value_type())?;
py_list.append(pyo3::types::PyTuple::new(py, [py_key, py_val])?)?;
}
Ok(py_list.into_any().unbind())
}
/// Convert a Fluss `ROW` to a Python dict keyed by field name — matching
/// pyarrow's `StructArray.to_pylist()` shape.
fn row_to_pydict(py: Python, row: &dyn InternalRow, row_type: &RowType) -> PyResult<Py<PyAny>> {
let dict = PyDict::new(py);
for (i, field) in row_type.fields().iter().enumerate() {
let py_val = datum_to_python_value(py, row, i, field.data_type())?;
dict.set_item(field.name(), py_val)?;
}
Ok(dict.into_any().unbind())
}
/// Convert element `i` of a binary `FlussArray` to a Python value. A binary
/// array needs the explicit nested type to decode MAP/ROW elements (the generic
/// `DataGetters` trait getters cover only scalars and nested arrays), so this
/// dispatches through `FlussArray`'s inherent typed getters.
fn array_elem_to_python(
py: Python,
arr: &FlussArray,
i: usize,
dt: &DataType,
) -> PyResult<Py<PyAny>> {
if arr.is_null_at(i) {
return Ok(py.None());
}
match dt {
DataType::Array(array_type) => {
let nested = arr
.get_array(i)
.map_err(|e| FlussError::from_core_error(&e))?;
array_to_pylist(py, &nested, array_type.get_element_type())
}
DataType::Map(map_type) => {
let nested = arr
.get_map(i, map_type.key_type(), map_type.value_type())
.map_err(|e| FlussError::from_core_error(&e))?;
map_to_pylist(py, &nested, map_type)
}
DataType::Row(row_type) => {
let nested = arr
.get_row(i, row_type)
.map_err(|e| FlussError::from_core_error(&e))?;
row_to_pydict(py, &nested, row_type)
}
_ => datum_to_python_value(py, arr, i, dt),
}
}
/// Convert Rust Decimal to Python decimal.Decimal
fn rust_decimal_to_python(py: Python, decimal: &fcore::row::Decimal) -> PyResult<Py<PyAny>> {
let decimal_ty = get_decimal_type(py)?;
let decimal_str = decimal.to_string();
let py_decimal = decimal_ty.call1((decimal_str,))?;
Ok(py_decimal.into_any().unbind())
}
/// Convert Rust Date (days since epoch) to Python datetime.date
fn rust_date_to_python(py: Python, date: fcore::row::Date) -> PyResult<Py<PyAny>> {
let days_since_epoch = date.get_inner();
let epoch = jiff::civil::date(1970, 1, 1);
let civil_date = epoch + jiff::Span::new().days(days_since_epoch as i64);
let py_date = PyDate::new(
py,
civil_date.year() as i32,
civil_date.month() as u8,
civil_date.day() as u8,
)?;
Ok(py_date.into_any().unbind())
}
/// Convert Rust Time (millis since midnight) to Python datetime.time
fn rust_time_to_python(py: Python, time: fcore::row::Time) -> PyResult<Py<PyAny>> {
let millis = time.get_inner() as i64;
let hours = millis / MILLIS_PER_HOUR;
let minutes = (millis % MILLIS_PER_HOUR) / MILLIS_PER_MINUTE;
let seconds = (millis % MILLIS_PER_MINUTE) / MILLIS_PER_SECOND;
let microseconds = (millis % MILLIS_PER_SECOND) * MICROS_PER_MILLI;
let py_time = PyTime::new(
py,
hours as u8,
minutes as u8,
seconds as u8,
microseconds as u32,
None,
)?;
Ok(py_time.into_any().unbind())
}
/// Convert Rust TimestampNtz to Python naive datetime
fn rust_timestamp_ntz_to_python(py: Python, ts: fcore::row::TimestampNtz) -> PyResult<Py<PyAny>> {
let millis = ts.get_millisecond();
let nanos = ts.get_nano_of_millisecond();
let total_micros = millis * MICROS_PER_MILLI + (nanos as i64 / NANOS_PER_MICRO);
// Convert to civil datetime via jiff
let timestamp = jiff::Timestamp::from_microsecond(total_micros)
.map_err(|e| FlussError::new_err(format!("Invalid timestamp: {e}")))?;
let civil_dt = timestamp.to_zoned(jiff::tz::TimeZone::UTC).datetime();
let py_dt = PyDateTime::new(
py,
civil_dt.year() as i32,
civil_dt.month() as u8,
civil_dt.day() as u8,
civil_dt.hour() as u8,
civil_dt.minute() as u8,
civil_dt.second() as u8,
(civil_dt.subsec_nanosecond() / 1000) as u32, // microseconds
None,
)?;
Ok(py_dt.into_any().unbind())
}
/// Convert Rust TimestampLtz to Python timezone-aware datetime (UTC)
fn rust_timestamp_ltz_to_python(py: Python, ts: fcore::row::TimestampLtz) -> PyResult<Py<PyAny>> {
let millis = ts.get_epoch_millisecond();
let nanos = ts.get_nano_of_millisecond();
let total_micros = millis * MICROS_PER_MILLI + (nanos as i64 / NANOS_PER_MICRO);
// Convert to civil datetime via jiff
let timestamp = jiff::Timestamp::from_microsecond(total_micros)
.map_err(|e| FlussError::new_err(format!("Invalid timestamp: {e}")))?;
let civil_dt = timestamp.to_zoned(jiff::tz::TimeZone::UTC).datetime();
let utc = get_utc_timezone(py)?;
let py_dt = PyDateTime::new(
py,
civil_dt.year() as i32,
civil_dt.month() as u8,
civil_dt.day() as u8,
civil_dt.hour() as u8,
civil_dt.minute() as u8,
civil_dt.second() as u8,
(civil_dt.subsec_nanosecond() / 1000) as u32, // microseconds
Some(&utc),
)?;
Ok(py_dt.into_any().unbind())
}
/// Convert an InternalRow to a Python dictionary
pub fn internal_row_to_dict(
py: Python,
row: &dyn fcore::row::InternalRow,
table_info: &fcore::metadata::TableInfo,
) -> PyResult<Py<PyAny>> {
let row_type = table_info.row_type();
let fields = row_type.fields();
let dict = PyDict::new(py);
for (pos, field) in fields.iter().enumerate() {
let value = datum_to_python_value(py, row, pos, field.data_type())?;
dict.set_item(field.name(), value)?;
}
Ok(dict.into_any().unbind())
}
/// Cached decimal.Decimal type
/// Uses PyOnceLock for thread-safety and subinterpreter compatibility.
static DECIMAL_TYPE: PyOnceLock<Py<PyType>> = PyOnceLock::new();
/// Cached UTC timezone
static UTC_TIMEZONE: PyOnceLock<Py<PyAny>> = PyOnceLock::new();
/// Cached UTC epoch type
static UTC_EPOCH: PyOnceLock<Py<PyAny>> = PyOnceLock::new();
/// Get the cached decimal.Decimal type, importing it once per interpreter.
fn get_decimal_type(py: Python) -> PyResult<Bound<PyType>> {
let ty = DECIMAL_TYPE.get_or_try_init(py, || -> PyResult<_> {
let decimal_mod = py.import("decimal")?;
let decimal_ty = decimal_mod.getattr("Decimal")?.downcast_into::<PyType>()?;
Ok(decimal_ty.unbind())
})?;
Ok(ty.bind(py).clone())
}
/// Get the cached UTC timezone (datetime.timezone.utc), creating it once per interpreter.
fn get_utc_timezone(py: Python) -> PyResult<Bound<PyTzInfo>> {
let tz = UTC_TIMEZONE.get_or_try_init(py, || -> PyResult<_> {
let datetime_mod = py.import("datetime")?;
let timezone = datetime_mod.getattr("timezone")?;
let utc = timezone.getattr("utc")?;
Ok(utc.unbind())
})?;
// Downcast to PyTzInfo for use with PyDateTime::new()
Ok(tz.bind(py).clone().downcast_into::<PyTzInfo>()?)
}
/// Get the cached UTC epoch datetime, creating it once per interpreter.
fn get_utc_epoch(py: Python) -> PyResult<Bound<PyAny>> {
let epoch = UTC_EPOCH.get_or_try_init(py, || -> PyResult<_> {
let datetime_mod = py.import("datetime")?;
let timezone = datetime_mod.getattr("timezone")?;
let utc = timezone.getattr("utc")?;
let epoch = datetime_mod
.getattr("datetime")?
.call1((1970, 1, 1, 0, 0, 0, 0, &utc))?;
Ok(epoch.unbind())
})?;
Ok(epoch.bind(py).clone())
}
/// Validate that value is a decimal.Decimal instance.
fn ensure_is_decimal(value: &Bound<PyAny>) -> PyResult<()> {
let decimal_ty = get_decimal_type(value.py())?;
if !value.is_instance(&decimal_ty.into_any())? {
return Err(FlussError::new_err(format!(
"Expected decimal.Decimal, got {}",
get_type_name(value)
)));
}
Ok(())
}
/// Convert Python decimal.Decimal to Datum::Decimal.
/// Only accepts decimal.Decimal
fn python_decimal_to_datum(
value: &Bound<PyAny>,
precision: u32,
scale: u32,
) -> PyResult<fcore::row::Datum<'static>> {
use std::str::FromStr;
ensure_is_decimal(value)?;
let decimal_str: String = value.str()?.extract()?;
let bd = bigdecimal::BigDecimal::from_str(&decimal_str).map_err(|e| {
FlussError::new_err(format!("Failed to parse decimal '{decimal_str}': {e}"))
})?;
let decimal = fcore::row::Decimal::from_big_decimal(bd, precision, scale).map_err(|e| {
FlussError::new_err(format!(
"Failed to convert decimal '{decimal_str}' to DECIMAL({precision}, {scale}): {e}"
))
})?;
Ok(fcore::row::Datum::Decimal(decimal))
}
/// Convert Python datetime.date to Datum::Date.
fn python_date_to_datum(value: &Bound<PyAny>) -> PyResult<fcore::row::Datum<'static>> {
// Reject datetime.datetime (subclass of date) - use timestamp columns for those
if value.downcast::<PyDateTime>().is_ok() {
return Err(FlussError::new_err(
"Expected datetime.date, got datetime.datetime. Use a TIMESTAMP column for datetime values.",
));
}
let date = value.downcast::<PyDate>().map_err(|_| {
FlussError::new_err(format!(
"Expected datetime.date, got {}",
get_type_name(value)
))
})?;
let year = date.get_year();
let month = date.get_month();
let day = date.get_day();
// Calculate days since Unix epoch (1970-01-01)
let civil_date = jiff::civil::date(year as i16, month as i8, day as i8);
let epoch = jiff::civil::date(1970, 1, 1);
let days_since_epoch = (civil_date - epoch).get_days();
Ok(fcore::row::Datum::Date(fcore::row::Date::new(
days_since_epoch,
)))
}
/// Convert Python datetime.time to Datum::Time.
/// Uses PyO3's native PyTime type for efficient access.
///
/// Note: Fluss TIME is always stored as milliseconds since midnight (i32) regardless
/// of the schema's precision setting. This matches the Java Fluss wire protocol.
/// Sub-millisecond precision (microseconds not divisible by 1000) will raise an error
/// to prevent silent data loss and ensure fail-fast behavior.
fn python_time_to_datum(value: &Bound<PyAny>) -> PyResult<fcore::row::Datum<'static>> {
let time = value.downcast::<PyTime>().map_err(|_| {
FlussError::new_err(format!(
"Expected datetime.time, got {}",
get_type_name(value)
))
})?;
let hour = time.get_hour() as i32;
let minute = time.get_minute() as i32;
let second = time.get_second() as i32;
let microsecond = time.get_microsecond() as i32;
// Strict validation: reject sub-millisecond precision
if microsecond % MICROS_PER_MILLI as i32 != 0 {
return Err(FlussError::new_err(format!(
"TIME values with sub-millisecond precision are not supported. \
Got time with {microsecond} microseconds (not divisible by 1000). \
Fluss stores TIME as milliseconds since midnight. \
Please round to milliseconds before insertion."
)));
}
// Convert to milliseconds since midnight
let millis = hour * MILLIS_PER_HOUR as i32
+ minute * MILLIS_PER_MINUTE as i32
+ second * MILLIS_PER_SECOND as i32
+ microsecond / MICROS_PER_MILLI as i32;
Ok(fcore::row::Datum::Time(fcore::row::Time::new(millis)))
}
/// Convert Python datetime-like object to Datum::TimestampNtz.
/// Supports: datetime.datetime (naive preferred), pd.Timestamp, np.datetime64
fn python_datetime_to_timestamp_ntz(value: &Bound<PyAny>) -> PyResult<fcore::row::Datum<'static>> {
let (epoch_millis, nano_of_milli) = extract_datetime_components_ntz(value)?;
let ts = fcore::row::TimestampNtz::from_millis_nanos(epoch_millis, nano_of_milli)
.map_err(|e| FlussError::new_err(format!("Failed to create TimestampNtz: {e}")))?;
Ok(fcore::row::Datum::TimestampNtz(ts))
}
/// Convert Python datetime-like object to Datum::TimestampLtz.
/// For naive datetimes, assumes UTC. For aware datetimes, converts to UTC.
/// Supports: datetime.datetime, pd.Timestamp, np.datetime64
fn python_datetime_to_timestamp_ltz(value: &Bound<PyAny>) -> PyResult<fcore::row::Datum<'static>> {
let (epoch_millis, nano_of_milli) = extract_datetime_components_ltz(value)?;
let ts = fcore::row::TimestampLtz::from_millis_nanos(epoch_millis, nano_of_milli)
.map_err(|e| FlussError::new_err(format!("Failed to create TimestampLtz: {e}")))?;
Ok(fcore::row::Datum::TimestampLtz(ts))
}
/// Extract epoch milliseconds for TimestampNtz (wall-clock time, no timezone conversion).
/// Uses integer arithmetic to avoid float precision issues.
/// For clarity, tz-aware datetimes are rejected - use TimestampLtz for those.
fn extract_datetime_components_ntz(value: &Bound<PyAny>) -> PyResult<(i64, i32)> {
// Try PyDateTime first
if let Ok(dt) = value.downcast::<PyDateTime>() {
// Reject tz-aware datetime for NTZ - it's ambiguous what the user wants
let tzinfo = dt.getattr("tzinfo")?;
if !tzinfo.is_none() {
return Err(FlussError::new_err(
"TIMESTAMP (without timezone) requires a naive datetime. \
Got timezone-aware datetime. Either remove tzinfo or use TIMESTAMP_LTZ column.",
));
}
return datetime_to_epoch_millis_as_utc(dt);
}
// Check for pandas Timestamp by verifying module name
if is_pandas_timestamp(value) {
// For NTZ, reject tz-aware pandas Timestamps for consistency with datetime behavior
if let Ok(tz) = value.getattr("tz") {
if !tz.is_none() {
return Err(FlussError::new_err(
"TIMESTAMP (without timezone) requires a naive pd.Timestamp. \
Got timezone-aware Timestamp. Either use tz_localize(None) or use TIMESTAMP_LTZ column.",
));
}
}
// Naive pandas Timestamp: .value is nanoseconds since epoch (wall-clock as UTC)
let nanos: i64 = value.getattr("value")?.extract()?;
return Ok(nanos_to_millis_and_submillis(nanos));
}
// Try to_pydatetime() for objects that support it
if let Ok(py_dt) = value.call_method0("to_pydatetime") {
if let Ok(dt) = py_dt.downcast::<PyDateTime>() {
let tzinfo = dt.getattr("tzinfo")?;
if !tzinfo.is_none() {
return Err(FlussError::new_err(
"TIMESTAMP (without timezone) requires a naive datetime. \
Got timezone-aware value. Use TIMESTAMP_LTZ column instead.",
));
}
return datetime_to_epoch_millis_as_utc(dt);
}
}
Err(FlussError::new_err(format!(
"Expected naive datetime.datetime or pd.Timestamp, got {}",
get_type_name(value)
)))
}
/// Extract epoch milliseconds for TimestampLtz (instant in time, UTC-based).
/// For naive datetimes, assumes UTC. For aware datetimes, converts to UTC.
fn extract_datetime_components_ltz(value: &Bound<PyAny>) -> PyResult<(i64, i32)> {
// Try PyDateTime first
if let Ok(dt) = value.downcast::<PyDateTime>() {
// Check if timezone-aware
let tzinfo = dt.getattr("tzinfo")?;
if tzinfo.is_none() {
// Naive datetime: assume UTC (treat components as UTC time)
return datetime_to_epoch_millis_as_utc(dt);
} else {
// Aware datetime: use timedelta from epoch to get correct UTC instant
return datetime_to_epoch_millis_utc_aware(dt);
}
}
// Check for pandas Timestamp
if is_pandas_timestamp(value) {
// pandas Timestamp.value is always nanoseconds since UTC epoch
let nanos: i64 = value.getattr("value")?.extract()?;
return Ok(nanos_to_millis_and_submillis(nanos));
}
// Try to_pydatetime()
if let Ok(py_dt) = value.call_method0("to_pydatetime") {
if let Ok(dt) = py_dt.downcast::<PyDateTime>() {
let tzinfo = dt.getattr("tzinfo")?;
if tzinfo.is_none() {
return datetime_to_epoch_millis_as_utc(dt);
} else {
return datetime_to_epoch_millis_utc_aware(dt);
}
}
}
Err(FlussError::new_err(format!(
"Expected datetime.datetime or pd.Timestamp, got {}",
get_type_name(value)
)))
}
/// Convert datetime components to epoch milliseconds treating them as UTC
fn datetime_to_epoch_millis_as_utc(dt: &Bound<'_, PyDateTime>) -> PyResult<(i64, i32)> {
let year = dt.get_year();
let month = dt.get_month();
let day = dt.get_day();
let hour = dt.get_hour();
let minute = dt.get_minute();
let second = dt.get_second();
let microsecond = dt.get_microsecond();
// Create jiff civil datetime and convert to UTC timestamp
// Safe casts: hour (0-23), minute (0-59), second (0-59) all fit in i8
let civil_dt = jiff::civil::date(year as i16, month as i8, day as i8).at(
hour as i8,
minute as i8,
second as i8,
microsecond as i32 * 1000,
);
let timestamp = jiff::tz::Offset::UTC
.to_timestamp(civil_dt)
.map_err(|e| FlussError::new_err(format!("Invalid datetime: {e}")))?;
let millis = timestamp.as_millisecond();
let nano_of_milli = (timestamp.subsec_nanosecond() % NANOS_PER_MILLI as i32) as i32;
Ok((millis, nano_of_milli))
}
/// Convert timezone-aware datetime to epoch milliseconds using Python's timedelta.
/// This correctly handles timezone conversions by computing (dt - UTC_EPOCH).
/// The UTC epoch is cached for performance.
fn datetime_to_epoch_millis_utc_aware(dt: &Bound<'_, PyDateTime>) -> PyResult<(i64, i32)> {
let py = dt.py();
let epoch = get_utc_epoch(py)?;
// Compute delta = dt - epoch (this handles timezone conversion correctly)
let delta = dt.call_method1("__sub__", (epoch,))?;
let delta = delta.downcast::<PyDelta>()?;
// Extract components using integer arithmetic
let days = delta.get_days() as i64;
let seconds = delta.get_seconds() as i64;
let microseconds = delta.get_microseconds() as i64;
// Total milliseconds (note: days can be negative for dates before epoch)
let total_micros = days * MICROS_PER_DAY + seconds * MICROS_PER_SECOND + microseconds;
let millis = total_micros / MICROS_PER_MILLI;
let nano_of_milli = ((total_micros % MICROS_PER_MILLI) * MICROS_PER_MILLI) as i32;
// Handle negative microseconds remainder
let (millis, nano_of_milli) = if nano_of_milli < 0 {
(millis - 1, nano_of_milli + NANOS_PER_MILLI as i32)
} else {
(millis, nano_of_milli)
};
Ok((millis, nano_of_milli))
}
/// Convert nanoseconds to (milliseconds, nano_of_millisecond)
fn nanos_to_millis_and_submillis(nanos: i64) -> (i64, i32) {
let millis = nanos / NANOS_PER_MILLI;
let nano_of_milli = (nanos % NANOS_PER_MILLI) as i32;
// Handle negative nanoseconds correctly (Euclidean remainder)
if nano_of_milli < 0 {
(millis - 1, nano_of_milli + NANOS_PER_MILLI as i32)
} else {
(millis, nano_of_milli)
}
}
/// Check if value is a pandas Timestamp by examining its type.
fn is_pandas_timestamp(value: &Bound<PyAny>) -> bool {
// Check module and class name to avoid importing pandas
if let Ok(cls) = value.get_type().getattr("__module__") {
if let Ok(module) = cls.extract::<&str>() {
if module.starts_with("pandas") {
if let Ok(name) = value.get_type().getattr("__name__") {
if let Ok(name_str) = name.extract::<&str>() {
return name_str == "Timestamp";
}
}
}
}
}
false
}
/// Get type name
fn get_type_name(value: &Bound<PyAny>) -> String {
value
.get_type()
.name()
.map(|s| s.to_string())
.unwrap_or_else(|_| "unknown".to_string())
}
/// Thin Python iterator over [`fcore::client::SyncRecordBatchLogReader`].
/// Used internally as the backing iterator for
/// ``pa.RecordBatchReader.from_batches()``; not registered on the module.
#[pyclass]
struct PyRecordBatchLogReader {
sync_reader: fcore::client::SyncRecordBatchLogReader,
}
#[pymethods]
impl PyRecordBatchLogReader {
fn __iter__(slf: PyRef<'_, Self>) -> PyRef<'_, Self> {
slf
}
fn __next__(&mut self, py: Python) -> PyResult<Option<Py<PyAny>>> {
let result = py.detach(|| self.sync_reader.next().transpose());
match result {
Ok(Some(batch)) => {
let py_batch = batch
.to_pyarrow(py)
.map_err(|e| FlussError::new_err(format!("Failed to convert batch: {e}")))?;
Ok(Some(py_batch.unbind()))
}
Ok(None) => Ok(None),
Err(arrow_err) => Err(FlussError::new_err(format!(
"Error reading batch: {arrow_err}"
))),
}
}
}
/// Wraps the two scanner variants so we never have an impossible state
/// (both None or both Some).
enum ScannerKind {
Record(fcore::client::LogScanner),
Batch(fcore::client::RecordBatchLogScanner),
}
impl ScannerKind {
fn as_record(&self) -> PyResult<&fcore::client::LogScanner> {
match self {
Self::Record(s) => Ok(s),
Self::Batch(_) => Err(FlussError::new_err(
"poll() requires a record-based scanner. Use new_scan().create_log_scanner().",
)),
}
}
fn as_batch(&self) -> PyResult<&fcore::client::RecordBatchLogScanner> {
match self {
Self::Batch(s) => Ok(s),
Self::Record(_) => Err(FlussError::new_err(
"This method requires a batch-based scanner. Use new_scan().create_record_batch_log_scanner().",
)),
}
}
}
/// Dispatch a method call to whichever scanner variant is active.
/// Both `LogScanner` and `RecordBatchLogScanner` share the same subscribe interface.
macro_rules! with_scanner {
($scanner:expr, $method:ident($($arg:expr),*)) => {
match $scanner.as_ref() {
ScannerKind::Record(s) => s.$method($($arg),*).await,
ScannerKind::Batch(s) => s.$method($($arg),*).await,
}
};
}
/// Scanner for reading log data from a Fluss table.
///
/// This scanner supports two modes:
/// - Record-based scanning via `poll()` - returns individual records with metadata
/// - Batch-based scanning via `poll_arrow()` / `poll_record_batch()` - returns Arrow batches
#[pyclass]
pub struct LogScanner {
kind: Arc<ScannerKind>,
admin: Arc<fcore::client::FlussAdmin>,
table_info: fcore::metadata::TableInfo,
/// The projected Arrow schema to use for empty table creation
projected_schema: SchemaRef,
/// The projected row type to use for record-based scanning
projected_row_type: Arc<fcore::metadata::RowType>,
}
#[pymethods]
impl LogScanner {
/// Subscribe to a single bucket at a specific offset (non-partitioned tables).
///
/// Args:
/// bucket_id: The bucket ID to subscribe to
/// start_offset: The offset to start reading from (use EARLIEST_OFFSET for beginning)
fn subscribe(&self, py: Python, bucket_id: i32, start_offset: i64) -> PyResult<()> {
py.detach(|| {
TOKIO_RUNTIME.block_on(async {
with_scanner!(&self.kind, subscribe(bucket_id, start_offset))
.map_err(|e| FlussError::from_core_error(&e))
})
})
}
/// Subscribe to multiple buckets at specified offsets (non-partitioned tables).
///
/// Args:
/// bucket_offsets: A dict mapping bucket_id -> start_offset
fn subscribe_buckets(&self, py: Python, bucket_offsets: HashMap<i32, i64>) -> PyResult<()> {
py.detach(|| {
TOKIO_RUNTIME.block_on(async {
with_scanner!(&self.kind, subscribe_buckets(&bucket_offsets))
.map_err(|e| FlussError::from_core_error(&e))
})
})
}
/// Subscribe to a bucket within a specific partition (partitioned tables only).
///
/// Args:
/// partition_id: The partition ID (from PartitionInfo.partition_id)
/// bucket_id: The bucket ID within the partition
/// start_offset: The offset to start reading from (use EARLIEST_OFFSET for beginning)
fn subscribe_partition(
&self,
py: Python,
partition_id: i64,
bucket_id: i32,
start_offset: i64,
) -> PyResult<()> {
py.detach(|| {
TOKIO_RUNTIME.block_on(async {
with_scanner!(
&self.kind,
subscribe_partition(partition_id, bucket_id, start_offset)
)
.map_err(|e| FlussError::from_core_error(&e))
})
})
}
/// Subscribe to multiple partition+bucket combinations at once (partitioned tables only).
///
/// Args:
/// partition_bucket_offsets: A dict mapping (partition_id, bucket_id) tuples to start_offsets
fn subscribe_partition_buckets(
&self,
py: Python,
partition_bucket_offsets: HashMap<(i64, i32), i64>,
) -> PyResult<()> {
py.detach(|| {
TOKIO_RUNTIME.block_on(async {
with_scanner!(
&self.kind,
subscribe_partition_buckets(&partition_bucket_offsets)
)
.map_err(|e| FlussError::from_core_error(&e))
})
})
}
/// Unsubscribe from a specific bucket (non-partitioned tables only).
///
/// Args:
/// bucket_id: The bucket ID to unsubscribe from
fn unsubscribe(&self, py: Python, bucket_id: i32) -> PyResult<()> {
py.detach(|| {
TOKIO_RUNTIME.block_on(async {
with_scanner!(&self.kind, unsubscribe(bucket_id))
.map_err(|e| FlussError::from_core_error(&e))
})
})
}
/// Unsubscribe from a specific partition bucket (partitioned tables only).
///
/// Args:
/// partition_id: The partition ID to unsubscribe from
/// bucket_id: The bucket ID within the partition
fn unsubscribe_partition(&self, py: Python, partition_id: i64, bucket_id: i32) -> PyResult<()> {
py.detach(|| {
TOKIO_RUNTIME.block_on(async {
with_scanner!(&self.kind, unsubscribe_partition(partition_id, bucket_id))
.map_err(|e| FlussError::from_core_error(&e))
})
})
}
/// Poll for individual records with metadata.
///
/// Args:
/// timeout_ms: Timeout in milliseconds to wait for records
///
/// Returns:
/// ScanRecords grouped by bucket. Supports flat iteration
/// (`for rec in records`) and per-bucket access (`records.buckets()`,
/// `records.records(bucket)`, `records[bucket]`).
///
/// Note:
/// - Requires a record-based scanner (created with new_scan().create_log_scanner())
/// - Returns an empty ScanRecords if no records are available
/// - When timeout expires, returns an empty ScanRecords (NOT an error)
fn poll<'py>(&self, py: Python<'py>, timeout_ms: i64) -> PyResult<Bound<'py, PyAny>> {
if timeout_ms < 0 {
return Err(FlussError::new_err(format!(
"timeout_ms must be non-negative, got: {timeout_ms}"
)));
}
let timeout = Duration::from_millis(timeout_ms as u64);
let scanner = Arc::clone(&self.kind);
let projected_row_type = self.projected_row_type.clone();
future_into_py(py, async move {
let scan_records = scanner
.as_record()?
.poll(timeout)
.await
.map_err(|e| FlussError::from_core_error(&e))?;
Python::attach(|py| {
let mut records_by_bucket = IndexMap::new();
let mut total_count = 0usize;
for (bucket, records) in scan_records.into_records_by_buckets() {
let py_bucket = TableBucket::from_core(bucket);
let mut py_records = Vec::with_capacity(records.len());
for record in &records {
let scan_record = ScanRecord::from_core(py, record, &projected_row_type)?;
py_records.push(Py::new(py, scan_record)?);
total_count += 1;
}
records_by_bucket.insert(py_bucket, py_records);
}
Ok(ScanRecords {
records_by_bucket,
total_count,
})
})
})
}
/// Poll for batches with metadata.
///
/// Args:
/// timeout_ms: Timeout in milliseconds to wait for batches
///
/// Returns:
/// List of RecordBatch objects, each containing the Arrow batch along with
/// bucket, base_offset, and last_offset metadata.
///
/// Note:
/// - Requires a batch-based scanner (created with new_scan().create_record_batch_log_scanner())
/// - Returns an empty list if no batches are available
/// - When timeout expires, returns an empty list (NOT an error)
fn poll_record_batch<'py>(
&self,
py: Python<'py>,
timeout_ms: i64,
) -> PyResult<Bound<'py, PyAny>> {
if timeout_ms < 0 {
return Err(FlussError::new_err(format!(
"timeout_ms must be non-negative, got: {timeout_ms}"
)));
}
let timeout = Duration::from_millis(timeout_ms as u64);
let scanner = Arc::clone(&self.kind);
future_into_py(py, async move {
let scan_batches = scanner
.as_batch()?
.poll(timeout)
.await
.map_err(|e| FlussError::from_core_error(&e))?;
Python::attach(|py| {
scan_batches
.into_iter()
.map(|sb| Py::new(py, RecordBatch::from_scan_batch(sb)))
.collect::<PyResult<Vec<_>>>()
})
})
}
/// Poll for new records as an Arrow Table.
///
/// Args:
/// timeout_ms: Timeout in milliseconds to wait for records
///
/// Returns:
/// PyArrow Table containing the polled records (batches merged)
///
/// Note:
/// - Requires a batch-based scanner (created with new_scan().create_record_batch_log_scanner())
/// - Returns an empty table (with correct schema) if no records are available
/// - When timeout expires, returns an empty table (NOT an error)
fn poll_arrow<'py>(&self, py: Python<'py>, timeout_ms: i64) -> PyResult<Bound<'py, PyAny>> {
if timeout_ms < 0 {
return Err(FlussError::new_err(format!(
"timeout_ms must be non-negative, got: {timeout_ms}"
)));
}
let timeout = Duration::from_millis(timeout_ms as u64);
let scanner = Arc::clone(&self.kind);
let projected_schema = self.projected_schema.clone();
future_into_py(py, async move {
let scan_batches = scanner
.as_batch()?
.poll(timeout)
.await
.map_err(|e| FlussError::from_core_error(&e))?;
let arrow_batches = scan_batches
.into_iter()
.map(|sb| Arc::new(sb.into_batch()))
.collect();
Python::attach(|py| Self::batches_to_arrow_table(py, arrow_batches, &projected_schema))
})
}
/// Create a lazy Arrow RecordBatchReader that reads until latest offsets.
///
/// This is a **blocking / synchronous** API: construction queries the
/// server for latest offsets (via ``block_on``), and each
/// ``RecordBatchReader.__next__()`` call blocks the calling thread until
/// the next batch is available. It is suitable for Arrow interop
/// (feeding into DuckDB, Polars, etc.) but should not be used
/// from ``asyncio`` coroutines -- see issue #545 for a planned
/// asyncio-native streaming alternative.
/// TODO(#545): Add asyncio-native streaming counterpart.
///
/// Returns a PyArrow RecordBatchReader that lazily polls batches one at a
/// time. This is more memory-efficient than ``to_arrow()`` which loads all
/// data into a single table.
///
/// **Concurrency:** While this reader is alive, ``subscribe*`` and
/// ``unsubscribe*`` calls on the scanner are rejected with an error.
/// You should also avoid calling ``poll_arrow`` / ``poll_record_batch``
/// on the same scanner — these are not blocked by the guard, but they
/// share the underlying fetch buffer with the reader and would
/// interleave batches between both consumers. Drop the reader before
/// resuming any of these operations.
///
/// You must call subscribe(), subscribe_buckets(), subscribe_partition(),
/// or subscribe_partition_buckets() first.
///
/// Returns:
/// ``pyarrow.RecordBatchReader`` yielding ``RecordBatch`` objects
fn to_arrow_batch_reader(&self, py: Python) -> PyResult<Py<PyAny>> {
let scanner = self.kind.as_batch()?;
let sync_reader = py
.detach(|| {
TOKIO_RUNTIME.block_on(async {
let reader = fcore::client::RecordBatchLogReader::new_until_latest(
scanner.new_shared_handle(),
&self.admin,
)
.await?;
Ok::<_, fcore::error::Error>(
reader.to_record_batch_reader(TOKIO_RUNTIME.handle().clone()),
)
})
})
.map_err(|e| FlussError::from_core_error(&e))?;
let py_schema = sync_reader
.schema()
.to_pyarrow(py)
.map_err(|e| FlussError::new_err(format!("Failed to convert schema: {e}")))?;
let py_iter = Py::new(py, PyRecordBatchLogReader { sync_reader })?;
let pyarrow = py.import("pyarrow")?;
let batch_reader = pyarrow
.getattr("RecordBatchReader")?
.call_method1("from_batches", (py_schema, py_iter))?;
Ok(batch_reader.into())
}
/// Convert all data to Arrow Table.
///
/// Reads from currently subscribed buckets until reaching their latest offsets.
/// Works for both partitioned and non-partitioned tables.
///
/// Materializes batches in Rust (``RecordBatchLogReader::collect_all_batches``)
/// then builds one PyArrow table, avoiding per-batch Python iteration.
///
/// You must call subscribe(), subscribe_buckets(), subscribe_partition(), or subscribe_partition_buckets() first.
///
/// Returns:
/// PyArrow Table containing all data from subscribed buckets
fn to_arrow<'py>(&self, py: Python<'py>) -> PyResult<Bound<'py, PyAny>> {
future_into_py(
py,
Self::scan_to_arrow_table(
Arc::clone(&self.kind),
Arc::clone(&self.admin),
self.projected_schema.clone(),
),
)
}
/// Convert all data to Pandas DataFrame.
///
/// Reads from currently subscribed buckets until reaching their latest offsets.
/// Works for both partitioned and non-partitioned tables.
///
/// You must call subscribe(), subscribe_buckets(), subscribe_partition(), or subscribe_partition_buckets() first.
///
/// Returns:
/// Pandas DataFrame containing all data from subscribed buckets
fn to_pandas<'py>(&self, py: Python<'py>) -> PyResult<Bound<'py, PyAny>> {
let kind = Arc::clone(&self.kind);
let admin = Arc::clone(&self.admin);
let projected_schema = self.projected_schema.clone();
future_into_py(py, async move {
let table = Self::scan_to_arrow_table(kind, admin, projected_schema).await?;
Python::attach(|py| table.call_method0(py, "to_pandas"))
})
}
fn __aiter__<'py>(slf: PyRef<'py, Self>) -> PyResult<Bound<'py, PyAny>> {
let py = slf.py();
// Single lock for the generic async generator
static ASYNC_GEN_FN: PyOnceLock<Py<PyAny>> = PyOnceLock::new();
let gen_fn = ASYNC_GEN_FN.get_or_init(py, || {
let code = pyo3::ffi::c_str!(
r#"
async def _async_scan_generic(scanner, method_name, timeout_ms):
poll_method = getattr(scanner, method_name)
while True:
for item in await poll_method(timeout_ms):
yield item
"#
);
let globals = pyo3::types::PyDict::new(py);
py.run(code, Some(&globals), None).unwrap();
globals
.get_item("_async_scan_generic")
.unwrap()
.unwrap()
.unbind()
});
let method_name = match slf.kind.as_ref() {
ScannerKind::Record(_) => "poll",
ScannerKind::Batch(_) => "poll_record_batch",
};
gen_fn.bind(py).call1((
slf.into_bound_py_any(py)?,
method_name,
DEFAULT_POLL_INTERVAL_MS,
))
}
fn __repr__(&self) -> String {
format!("LogScanner(table={})", self.table_info.table_path)
}
}
impl LogScanner {
fn new(
scanner: ScannerKind,
admin: Arc<fcore::client::FlussAdmin>,
table_info: fcore::metadata::TableInfo,
projected_schema: SchemaRef,
projected_row_type: Arc<fcore::metadata::RowType>,
) -> Self {
Self {
kind: Arc::new(scanner),
admin,
table_info,
projected_schema,
projected_row_type,
}
}
/// Read until the latest offsets and build one PyArrow Table.
async fn scan_to_arrow_table(
kind: Arc<ScannerKind>,
admin: Arc<fcore::client::FlussAdmin>,
projected_schema: SchemaRef,
) -> PyResult<Py<PyAny>> {
let scanner = kind.as_batch()?;
let mut reader = fcore::client::RecordBatchLogReader::new_until_latest(
scanner.new_shared_handle(),
&admin,
)
.await
.map_err(|e| FlussError::from_core_error(&e))?;
let batches: Vec<Arc<ArrowRecordBatch>> = reader
.collect_all_batches()
.await
.map_err(|e| FlussError::from_core_error(&e))?
.into_iter()
.map(|sb| Arc::new(sb.into_batch()))
.collect();
Python::attach(|py| Self::batches_to_arrow_table(py, batches, &projected_schema))
}
/// Convert Arrow record batches to a PyArrow Table (or empty table if no batches).
fn batches_to_arrow_table(
py: Python<'_>,
batches: Vec<Arc<ArrowRecordBatch>>,
projected_schema: &SchemaRef,
) -> PyResult<Py<PyAny>> {
if batches.is_empty() {
let py_schema = projected_schema
.as_ref()
.to_pyarrow(py)
.map_err(|e| FlussError::new_err(format!("Failed to convert schema: {e}")))?;
let pyarrow = py.import("pyarrow")?;
let empty_table = pyarrow
.getattr("Table")?
.call_method1("from_batches", (vec![] as Vec<Py<PyAny>>, py_schema))?;
Ok(empty_table.into())
} else {
Utils::combine_batches_to_table(py, batches)
}
}
}
/// One-shot bounded scanner over a single bucket.
///
/// Obtained via `table.new_scan().limit(n).create_bucket_batch_scanner(bucket)`.
/// The scan runs on the first `next_batch()` and yields its single batch once,
/// then is spent. Honors the configured limit and any projection.
#[pyclass]
pub struct BatchScanner {
inner: Arc<Mutex<LimitBatchScanner>>,
bucket: TableBucket,
projected_schema: SchemaRef,
}
#[pymethods]
impl BatchScanner {
/// The bucket scanned by this batch scanner.
#[getter]
fn bucket(&self) -> TableBucket {
self.bucket.clone()
}
/// Run the scan and return its batch, or `None` once the scanner is spent.
///
/// The scan runs on the first call and is not retried; on error, create a
/// new scanner.
fn next_batch<'py>(&self, py: Python<'py>) -> PyResult<Bound<'py, PyAny>> {
let inner = Arc::clone(&self.inner);
future_into_py(py, async move {
let mut scanner = inner.lock().await;
let batch = scanner
.next_batch()
.await
.map_err(|e| FlussError::from_core_error(&e))?;
Python::attach(|py| match batch {
Some(sb) => Ok(Some(Py::new(py, RecordBatch::from_scan_batch(sb))?)),
None => Ok(None),
})
})
}
/// Drain the scanner into all of its batches.
fn collect_all_batches<'py>(&self, py: Python<'py>) -> PyResult<Bound<'py, PyAny>> {
let inner = Arc::clone(&self.inner);
future_into_py(py, async move {
let mut scanner = inner.lock().await;
let batches = scanner
.collect_all_batches()
.await
.map_err(|e| FlussError::from_core_error(&e))?;
Python::attach(|py| {
batches
.into_iter()
.map(|sb| Py::new(py, RecordBatch::from_scan_batch(sb)))
.collect::<PyResult<Vec<_>>>()
})
})
}
/// Drain the scanner into a PyArrow Table (empty, with the projected schema,
/// when the scan yields nothing).
fn to_arrow<'py>(&self, py: Python<'py>) -> PyResult<Bound<'py, PyAny>> {
future_into_py(
py,
Self::scan_to_arrow_table(Arc::clone(&self.inner), self.projected_schema.clone()),
)
}
/// Drain the scanner into a Pandas DataFrame.
fn to_pandas<'py>(&self, py: Python<'py>) -> PyResult<Bound<'py, PyAny>> {
let inner = Arc::clone(&self.inner);
let projected_schema = self.projected_schema.clone();
future_into_py(py, async move {
let table = Self::scan_to_arrow_table(inner, projected_schema).await?;
Python::attach(|py| table.call_method0(py, "to_pandas"))
})
}
fn __repr__(&self) -> String {
format!("BatchScanner(bucket={})", self.bucket.__str__())
}
}
impl BatchScanner {
fn new(scanner: LimitBatchScanner, bucket: TableBucket, projected_schema: SchemaRef) -> Self {
Self {
inner: Arc::new(Mutex::new(scanner)),
bucket,
projected_schema,
}
}
/// Drain the scanner into one PyArrow Table.
async fn scan_to_arrow_table(
inner: Arc<Mutex<LimitBatchScanner>>,
projected_schema: SchemaRef,
) -> PyResult<Py<PyAny>> {
let mut scanner = inner.lock().await;
let batches = scanner
.collect_all_batches()
.await
.map_err(|e| FlussError::from_core_error(&e))?
.into_iter()
.map(|sb| Arc::new(sb.into_batch()))
.collect();
Python::attach(|py| LogScanner::batches_to_arrow_table(py, batches, &projected_schema))
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_nanos_to_millis_and_submillis() {
// Simple positive case
assert_eq!(nanos_to_millis_and_submillis(1_500_000), (1, 500_000));
// Exact millisecond boundary
assert_eq!(nanos_to_millis_and_submillis(2_000_000), (2, 0));
// Zero
assert_eq!(nanos_to_millis_and_submillis(0), (0, 0));
// Large value
assert_eq!(
nanos_to_millis_and_submillis(86_400_000_000_000), // 1 day in nanos
(86_400_000, 0)
);
// Negative: -1.5 milliseconds should be (-2 millis, +500_000 nanos)
// Because -1_500_000 nanos = -2ms + 500_000ns
assert_eq!(nanos_to_millis_and_submillis(-1_500_000), (-2, 500_000));
// Negative exact boundary
assert_eq!(nanos_to_millis_and_submillis(-2_000_000), (-2, 0));
// Small negative
assert_eq!(nanos_to_millis_and_submillis(-1), (-1, 999_999));
// Negative with sub-millisecond part
assert_eq!(nanos_to_millis_and_submillis(-500_000), (-1, 500_000));
}
}