Python topic publisher (#19)

feat: publish and query support for Python API

- Added topic publishing to Python bindings.
- cleaned up Python typing and documentation.
- added more Python API documentation
- changed cargo deny to ignore private crates

.gitignore

@@ -3,5 +3,7 @@
 # Rust
 target
 
-# pyenv overrides
+# Python
+*.so
+__pycache__
 .python-version

CONTRIBUTING.md

@@ -25,3 +25,17 @@ This should:
 - Create a tag for the new version
 
 Once the new tag is pushed to the repo, a new Github release and Docker image will be generated automatically by Github Actions.
+
+## Generating Python documentation
+
+Generating Python documentation requires pdoc:
+
+```shell
+pip install pdoc
+```
+
+To view the documentation, run:
+
+```shell
+pdoc ella -d google
+```

Cargo.toml

@@ -1,5 +1,4 @@
 [workspace]
-exclude = ["ella-cli"]
 members = [
     "ella-common",
     "ella-tensor",
@@ -9,6 +8,7 @@ members = [
     "ella-cli",
     "ella",
     "pyella",
+    "pyella/generate_typing",
 ]
 resolver = "2"
 

deny.toml

@@ -171,7 +171,7 @@ exceptions = [
 # published to private registries.
 # To see how to mark a crate as unpublished (to the official registry),
 # visit https://doc.rust-lang.org/cargo/reference/manifest.html#the-publish-field.
-ignore = false
+ignore = true
 # One or more private registries that you might publish crates to, if a crate
 # is only published to private registries, and ignore is true, the crate will
 # not have its license(s) checked

ella-engine/src/lazy.rs

@@ -43,6 +43,22 @@ impl Lazy {
         Ok(self.stream().await?.rows())
     }
 
+    pub fn limit(mut self, limit: usize) -> crate::Result<Self> {
+        self.plan = self.plan.try_map(|plan| {
+            LogicalPlanBuilder::from(plan)
+                .limit(0, Some(limit))?
+                .build()
+        })?;
+        Ok(self)
+    }
+
+    pub fn skip(mut self, skip: usize) -> crate::Result<Self> {
+        self.plan = self
+            .plan
+            .try_map(|plan| LogicalPlanBuilder::from(plan).limit(skip, None)?.build())?;
+        Ok(self)
+    }
+
     pub fn col<T, S>(&self, col: &str) -> crate::Result<Column<T, S>>
     where
         T: TensorValue,

ella-engine/src/plan.rs

@@ -82,6 +82,18 @@ impl Plan {
         };
         self
     }
+
+    pub fn try_map<F, E>(mut self, f: F) -> crate::Result<Self>
+    where
+        F: FnOnce(LogicalPlan) -> Result<LogicalPlan, E>,
+        crate::Error: From<E>,
+    {
+        self.inner = match self.inner {
+            PlanInner::Resolved(plan) => PlanInner::Resolved(f(plan)?),
+            PlanInner::Stub(plan) => PlanInner::Stub(f(plan)?),
+        };
+        Ok(self)
+    }
 }
 
 #[derive(Debug, Clone, PartialEq, Eq)]

ella-engine/src/registry/id.rs

@@ -571,7 +571,7 @@ fn encode_uuid_to_path(uuid: Uuid, root: &Path, prefix: Option<&str>, ext: &str)
 }
 
 // TODO: replace UUID logic once v7 is stabilized
-// See: [BlackrockNeurotech/ella#11](https://github.com/BlackrockNeurotech/ella/issues/11)
+// See: [CerebusOSS/ella#11](https://github.com/CerebusOSS/ella/issues/11)
 // Current implementation is taken from here: https://github.com/uuid-rs/uuid/blob/main/src/timestamp.rs
 
 fn new_uuid_now() -> Uuid {

ella-engine/src/table/topic/channel.rs

@@ -109,12 +109,7 @@ impl Clone for PublisherInner {
 
 impl Drop for PublisherInner {
     fn drop(&mut self) {
-        if self.is_active {
-            let active = self.active.fetch_sub(1, Ordering::Release) - 1;
-            if active == 0 {
-                self.stop.notify_one();
-            }
-        }
+        self.deactivate();
     }
 }
 
@@ -132,6 +127,16 @@ impl PublisherInner {
             is_active,
         }
     }
+
+    fn deactivate(&mut self) {
+        if self.is_active {
+            self.is_active = false;
+            let active = self.active.fetch_sub(1, Ordering::Release) - 1;
+            if active == 0 {
+                self.stop.notify_one();
+            }
+        }
+    }
 }
 
 #[derive(Debug, Clone)]
@@ -170,12 +175,13 @@ impl Sink<RecordBatch> for Publisher {
         self.inner.rw.poll_flush_unpin(cx)
     }
 
-    #[inline]
     fn poll_close(
         mut self: std::pin::Pin<&mut Self>,
         cx: &mut std::task::Context<'_>,
     ) -> Poll<std::result::Result<(), Self::Error>> {
-        self.inner.rw.poll_close_unpin(cx)
+        let res = futures::ready!(self.inner.rw.poll_close_unpin(cx));
+        self.inner.deactivate();
+        Poll::Ready(res)
     }
 }
 

ella-tensor/src/column.rs

@@ -89,7 +89,7 @@ pub trait Column: Debug + Send + Sync {
 
     fn row_shape(&self) -> Option<Dyn> {
         let shape = self.shape();
-        if shape.ndim() > 0 {
+        if shape.ndim() > 1 {
             Some(shape.remove_axis(Axis(0)))
         } else {
             None

ella-tensor/src/frame/data_frame.rs

@@ -1,9 +1,12 @@
 use std::{fmt::Display, ops::Deref, sync::Arc};
 
-use arrow::record_batch::RecordBatch;
+use arrow::{
+    datatypes::{Fields, Schema},
+    record_batch::RecordBatch,
+};
 use ella_common::row::RowFormat;
 
-use crate::{NamedColumn, Shape, Tensor, TensorValue};
+use crate::{tensor_schema, NamedColumn, Shape, Tensor, TensorValue};
 
 use super::{batch_to_columns, frame_to_batch, print::print_frames, Frame};
 
@@ -40,22 +43,52 @@ impl DataFrame {
         R::view(batch.num_rows(), &batch.schema().fields, batch.columns())
     }
 
+    pub fn ncols(&self) -> usize {
+        self.columns.len()
+    }
+
+    pub fn nrows(&self) -> usize {
+        self.rows
+    }
+
+    pub fn column(&self, i: usize) -> &NamedColumn {
+        &self.columns[i]
+    }
+
     pub fn pretty_print(&self) -> impl Display + '_ {
         print_frames(&[self])
     }
+
+    pub fn arrow_schema(&self) -> Schema {
+        Schema::new(
+            self.columns()
+                .map(|col| {
+                    Arc::new(tensor_schema(
+                        col.name().to_string(),
+                        col.tensor_type(),
+                        col.row_shape(),
+                        col.nullable(),
+                    ))
+                })
+                .collect::<Fields>(),
+        )
+    }
 }
 
 impl Frame for DataFrame {
+    #[inline]
     fn ncols(&self) -> usize {
-        self.columns.len()
+        self.ncols()
     }
 
+    #[inline]
     fn nrows(&self) -> usize {
-        self.rows
+        self.nrows()
     }
 
+    #[inline]
     fn column(&self, i: usize) -> &NamedColumn {
-        &self.columns[i]
+        self.column(i)
     }
 }
 

ella/src/table/publisher.rs

@@ -24,6 +24,10 @@ impl Publisher {
         let schema = self.arrow_schema.clone();
         RowSink::try_new(self, schema, buffer)
     }
+
+    pub fn arrow_schema(&self) -> &SchemaRef {
+        &self.arrow_schema
+    }
 }
 
 impl Debug for Publisher {

justfile

@@ -10,3 +10,6 @@ generate_release_notes version:
 
 generate_readme:
     cd ella; cargo rdme
+
+generate_python:
+    cargo run -p generate_typing --bin generate_data_types

pyella/Cargo.toml

@@ -25,6 +25,7 @@ thiserror = { workspace = true }
 once_cell = { workspace = true }
 serde = { workspace = true }
 serde_json = { workspace = true }
+tracing-subscriber = { workspace = true }
 
 mimalloc = { workspace = true, optional = true }
 

pyella/ella/__init__.py

@@ -1,21 +1,47 @@
 # ruff: noqa: E402, F403
 __all__ = [
+    "types",
+    "table",
+    "frame",
     "open",
     "connect",
     "column",
     "topic",
     "Ella",
-    "data_types",
+    "now",
+    "bool_",
+    "int8",
+    "int16",
+    "int32",
+    "int64",
+    "uint8",
+    "uint16",
+    "uint32",
+    "uint64",
+    "float32",
+    "float64",
+    "timestamp",
+    "duration",
+    "string",
 ]
 
-from maturin import import_hook as __import_hook
-__import_hook.install()
+from . import types, table, frame
 
-def __add_submodule(path, src):
-    import sys
-    sys.modules[path] = src
+from .types import (
+    bool_,
+    int8,
+    int16,
+    int32,
+    int64,
+    uint8,
+    uint16,
+    uint32,
+    uint64,
+    float32,
+    float64,
+    timestamp,
+    duration,
+    string,
+)
 
-from ella._internal import open, connect, column, topic, Ella, data_types
-from ella._internal.type_defs import *
-
-__add_submodule("ella.data_types", data_types)
+from ella._internal import open, connect, column, topic, Ella, now