CC Self-review. Applied lint. Using new method in the high level API (XarrayContext). This leads to two failed tests, but this could be caused by test errors.

Author: alxmrs

.gitignore

@@ -8,4 +8,4 @@ __pycache__
 .idea
 target
 test_data
-*.so
+*.so

Cargo.lock

@@ -7,14 +7,10 @@ exclude = ["perf_tests/*"]
 
 [dependencies]
 arrow = { version = "55.2.0", features = ["pyarrow"] }
-arrow-array = { version = "55.2.0" }
-arrow-schema = { version = "55.2.0" }
-async-trait = "0.1.89"
 datafusion = { version = "49.0.0" }
-datafusion-ffi = {version = "49.0.0" }
-futures = "0.3.31"
+datafusion-ffi = { version = "49.0.0" }
 pyo3 = { version = "0.24.1", features = ["extension-module"] }
-tokio = { version = "1.46.1", features = ["rt", "rt-multi-thread", "macros"] }
+tokio = { version = "1.46.1", features = ["rt"] }
 
 
 [build-dependencies]

Cargo.toml

@@ -73,10 +73,25 @@ impl PartitionStream for PyArrowStreamPartition {
                 Python::with_gil(|py| {
                     let bound = obj.bind(py);
 
-                    // Create a stream reader and collect all batches
                     match ArrowArrayStreamReader::from_pyarrow_bound(bound) {
-                        Ok(reader) => reader.filter_map(|result| result.ok()).collect(),
-                        Err(_) => vec![],
+                        Ok(reader) => {
+                            // Collect batches, propagating errors as warnings
+                            // In streaming context, we can't easily return errors,
+                            // so we log and skip failed batches
+                            reader
+                                .filter_map(|result| match result {
+                                    Ok(batch) => Some(batch),
+                                    Err(e) => {
+                                        eprintln!("Warning: Failed to read batch: {e}");
+                                        None
+                                    }
+                                })
+                                .collect()
+                        }
+                        Err(e) => {
+                            eprintln!("Warning: Failed to create stream reader: {e}");
+                            vec![]
+                        }
                     }
                 })
             }
@@ -86,7 +101,6 @@ impl PartitionStream for PyArrowStreamPartition {
             }
         };
 
-        // Wrap the batches in a MemoryStream
         Box::pin(
             MemoryStream::try_new(batches, Arc::clone(&self.schema), None)
                 .expect("MemoryStream creation should not fail with valid schema"),
@@ -190,28 +204,10 @@ impl LazyArrowStreamTable {
         )
     }
 
-    /// Get the schema of the table.
+    /// Get the schema of the table as a PyArrow Schema.
     fn schema(&self, py: Python<'_>) -> PyResult<PyObject> {
-        let schema = self.table.schema();
-        // Convert to PyArrow schema
-        let pyarrow = py.import("pyarrow")?;
-        let schema_cls = pyarrow.getattr("schema")?;
-
-        // Build field list
-        let fields: Vec<_> = schema
-            .fields()
-            .iter()
-            .map(|f: &arrow::datatypes::FieldRef| {
-                let field_fn = pyarrow.getattr("field").unwrap();
-                let dtype = arrow_type_to_pyarrow(py, f.data_type()).unwrap();
-                field_fn
-                    .call1((f.name(), dtype, f.is_nullable()))
-                    .unwrap()
-            })
-            .collect();
-
-        let result = schema_cls.call1((fields,))?;
-        Ok(result.into())
+        use arrow::pyarrow::ToPyArrow;
+        self.table.schema().to_pyarrow(py)
     }
 
     fn __repr__(&self) -> String {
@@ -222,84 +218,30 @@ impl LazyArrowStreamTable {
     }
 }
 
-/// Get schema from a Python object that has a schema attribute or method.
+/// Get schema from a Python object that has a schema attribute.
 ///
-/// This function extracts the schema WITHOUT consuming the stream, which is
+/// This extracts the schema WITHOUT consuming the stream, which is
 /// important for lazy evaluation.
 fn get_schema_from_stream(stream: &Bound<'_, PyAny>) -> PyResult<SchemaRef> {
     use arrow::datatypes::Schema;
     use arrow::pyarrow::FromPyArrow;
 
-    // Try to get the schema attribute (PyArrow RecordBatchReader has .schema property)
     let py_schema = stream.getattr("schema").map_err(|e| {
         pyo3::exceptions::PyTypeError::new_err(format!(
-            "Object does not have a 'schema' attribute: {e}. \
-             Expected a RecordBatchReader or similar object with a schema property."
+            "Object must have a 'schema' attribute (e.g., RecordBatchReader): {e}"
         ))
     })?;
 
-    // Convert the PyArrow Schema to Rust Schema
     let schema = Schema::from_pyarrow_bound(&py_schema).map_err(|e| {
-        pyo3::exceptions::PyTypeError::new_err(format!(
-            "Failed to convert schema: {e}"
-        ))
+        pyo3::exceptions::PyTypeError::new_err(format!("Failed to convert schema: {e}"))
     })?;
 
     Ok(Arc::new(schema))
 }
 
-/// Convert Arrow DataType to PyArrow type
-fn arrow_type_to_pyarrow(py: Python<'_>, dtype: &arrow::datatypes::DataType) -> PyResult<PyObject> {
-    let pyarrow = py.import("pyarrow")?;
-
-    use arrow::datatypes::DataType;
-    let result = match dtype {
-        DataType::Null => pyarrow.call_method0("null")?,
-        DataType::Boolean => pyarrow.call_method0("bool_")?,
-        DataType::Int8 => pyarrow.call_method0("int8")?,
-        DataType::Int16 => pyarrow.call_method0("int16")?,
-        DataType::Int32 => pyarrow.call_method0("int32")?,
-        DataType::Int64 => pyarrow.call_method0("int64")?,
-        DataType::UInt8 => pyarrow.call_method0("uint8")?,
-        DataType::UInt16 => pyarrow.call_method0("uint16")?,
-        DataType::UInt32 => pyarrow.call_method0("uint32")?,
-        DataType::UInt64 => pyarrow.call_method0("uint64")?,
-        DataType::Float16 => pyarrow.call_method0("float16")?,
-        DataType::Float32 => pyarrow.call_method0("float32")?,
-        DataType::Float64 => pyarrow.call_method0("float64")?,
-        DataType::Utf8 => pyarrow.call_method0("utf8")?,
-        DataType::LargeUtf8 => pyarrow.call_method0("large_utf8")?,
-        DataType::Binary => pyarrow.call_method0("binary")?,
-        DataType::LargeBinary => pyarrow.call_method0("large_binary")?,
-        DataType::Date32 => pyarrow.call_method0("date32")?,
-        DataType::Date64 => pyarrow.call_method0("date64")?,
-        DataType::Timestamp(unit, tz) => {
-            let unit_str = match unit {
-                arrow::datatypes::TimeUnit::Second => "s",
-                arrow::datatypes::TimeUnit::Millisecond => "ms",
-                arrow::datatypes::TimeUnit::Microsecond => "us",
-                arrow::datatypes::TimeUnit::Nanosecond => "ns",
-            };
-            match tz {
-                Some(tz) => pyarrow.call_method1("timestamp", (unit_str, tz.to_string()))?,
-                None => pyarrow.call_method1("timestamp", (unit_str,))?,
-            }
-        }
-        _ => {
-            // Fallback: convert to string and let PyArrow parse it
-            let type_str = format!("{dtype}");
-            pyarrow
-                .getattr("type_for_alias")?
-                .call1((type_str.as_str(),))?
-        }
-    };
-
-    Ok(result.into())
-}
-
 /// Python module initialization
 #[pymodule]
 fn _native(m: &Bound<'_, PyModule>) -> PyResult<()> {
     m.add_class::<LazyArrowStreamTable>()?;
     Ok(())
-}
+}

src/lib.rs

@@ -1,4 +1,16 @@
 from .df import read_xarray, from_map
-from .reader import XarrayRecordBatchReader, read_xarray_lazy
+from .reader import XarrayRecordBatchReader, read_xarray_lazy, read_xarray_table
 from .sql import XarrayContext
 from ._native import LazyArrowStreamTable
+
+__all__ = [
+    # High-level API (recommended)
+    "read_xarray_table",
+    "XarrayContext",
+    # Lower-level building blocks
+    "read_xarray",
+    "read_xarray_lazy",
+    "from_map",
+    "XarrayRecordBatchReader",
+    "LazyArrowStreamTable",
+]