Repository: mwlon/pcodec (language: python)

Issue:
Add missing options to `PyChunkConfig`
Several options and properties are missing from `PyChunkConfig`. This includes:
- The `float_quant_spec` property.
- The `int_mult_spec` `Provided` option.
- The `float_mult_spec` `Provided` option.

I think updating `PyChunkConfig` to include all `ChunkConfig` options should be relatively simple. Should we do it?


Discussion / hints:
Yes, we should do it. Float quant will add Enabled in Rust land soon, but once all these things are in, we should release 0.2.0 (since it extends the format).
Started work on this. Will submit PR soon.

Fix patch:
diff --git a/pco_python/src/lib.rs b/pco_python/src/lib.rs
index bf61560c..552ab310 100644
--- a/pco_python/src/lib.rs
+++ b/pco_python/src/lib.rs
@@ -1,7 +1,7 @@
 use half::f16;
 use numpy::PyArrayDyn;
 use pco::data_types::CoreDataType;
-use pco::{ChunkConfig, FloatMultSpec, IntMultSpec, PagingSpec, Progress};
+use pco::{ChunkConfig, FloatMultSpec, FloatQuantSpec, IntMultSpec, PagingSpec, Progress};
 use pyo3::exceptions::PyRuntimeError;
 use pyo3::prelude::{pymodule, FromPyObject, PyModule, PyResult, Python};
 use pyo3::{py_run, pyclass, pymethods, PyErr};
@@ -50,6 +50,85 @@ pub fn pco_err_to_py(pco: PcoError) -> PyErr {
   PyRuntimeError::new_err(format!("pco error: {}", pco))
 }
 
+#[pyclass(name = "IntMultSpec")]
+#[derive(Clone, Default)]
+pub struct PyIntMultSpec(IntMultSpec);
+
+/// Specifies if pcodec should use an integer multiplier to compress data.
+#[pymethods]
+impl PyIntMultSpec {
+  /// :returns: a IntMultSpec disabling integer multiplier detection.
+  #[staticmethod]
+  fn disabled() -> Self {
+    Self(IntMultSpec::Disabled)
+  }
+
+  /// :returns: a IntMultSpec enabling integer multiplier detection. This
+  /// will automatically try to detect whether to use the mode and which `base`
+  /// to use.
+  #[staticmethod]
+  fn enabled() -> Self {
+    Self(IntMultSpec::Enabled)
+  }
+
+  /// :returns: a IntMultSpec with a specific `base` to use for compression.
+  #[staticmethod]
+  fn provided(base: u64) -> Self {
+    Self(IntMultSpec::Provided(base))
+  }
+}
+
+#[pyclass(name = "FloatMultSpec")]
+#[derive(Clone, Default)]
+pub struct PyFloatMultSpec(FloatMultSpec);
+
+/// Specifies if pcodec should use a floating point multiplier to compress
+/// data.
+#[pymethods]
+impl PyFloatMultSpec {
+  /// :returns: a FloatMultSpec disabling floating point multiplier detection.
+  #[staticmethod]
+  fn disabled() -> Self {
+    Self(FloatMultSpec::Disabled)
+  }
+
+  /// :returns: a FloatMultSpec enabling floating point multiplier detection.
+  /// This will automatically try to detect whether to use the mode and which
+  /// `base` to use.
+  #[staticmethod]
+  fn enabled() -> Self {
+    Self(FloatMultSpec::Enabled)
+  }
+
+  /// :returns: a FloatMultSpec with a specific `base` to use for compression.
+  #[staticmethod]
+  fn provided(base: f64) -> Self {
+    Self(FloatMultSpec::Provided(base))
+  }
+}
+
+#[pyclass(name = "FloatQuantSpec")]
+#[derive(Clone, Default)]
+pub struct PyFloatQuantSpec(FloatQuantSpec);
+
+/// Specifies if pcodec should use floating point quantization to compress
+/// data.
+#[pymethods]
+impl PyFloatQuantSpec {
+  /// :returns: a FloatQuantSpec disabling floating point quantization.
+  #[staticmethod]
+  fn disabled() -> Self {
+    Self(FloatQuantSpec::Disabled)
+  }
+
+  /// :returns: a FloatQuantSpec specifying that pcodec should use quantization
+  /// with a specific number of `bits``.
+  #[staticmethod]
+  fn provided(bits: u32) -> Self {
+    Self(FloatQuantSpec::Provided(bits))
+  }
+}
+
 #[pyclass(name = "PagingSpec")]
 #[derive(Clone, Default)]
 pub struct PyPagingSpec(PagingSpec);
@@ -78,8 +157,9 @@ impl PyPagingSpec {
 pub struct PyChunkConfig {
   compression_level: usize,
   delta_encoding_order: Option<usize>,
-  int_mult_spec: String,
-  float_mult_spec: String,
+  int_mult_spec: PyIntMultSpec,
+  float_mult_spec: PyFloatMultSpec,
+  float_quant_spec: PyFloatQuantSpec,
   paging_spec: PyPagingSpec,
 }
 
@@ -89,15 +169,44 @@ impl PyChunkConfig {
   ///
   /// :param compression_level: a compression level from 0-12, where 12 takes
   /// the longest and compresses the most.
+  ///
   /// :param delta_encoding_order: either a delta encoding level from 0-7 or
   /// None. If set to None, pcodec will try to infer the optimal delta encoding
   /// order.
-  /// :param int_mult_spec: either 'enabled' or 'disabled'. If enabled, pcodec
-  /// will consider using int mult mode, which can substantially improve
-  /// compression ratio but decrease speed in some cases for integer types.
-  /// :param float_mult_spec: either 'enabled' or 'disabled'. If enabled, pcodec
-  /// will consider using float mult mode, which can substantially improve
-  /// compression ratio but decrease speed in some cases for float types.
+  ///
+  /// :param int_mult_spec: a IntMultSpec that configures whether integer
+  /// multiplier detection is enabled.
+  ///
+  /// Examples where this helps:
+  /// * nanosecond-precision timestamps that are mostly whole numbers of
+  /// microseconds, with a few exceptions
+  /// * integers `[7, 107, 207, 307, ... 100007]` shuffled
+  ///
+  /// When this is helpful, compression and decompression speeds can be
+  /// substantially reduced. This configuration may hurt
+  /// compression speed slightly even when it isn't helpful.
+  /// However, the compression ratio improvements tend to be large.
+  ///
+  /// :param float_mult_spec: a FloatMultSpec that configures whether float
+  /// multiplier detection is enabled.
+  ///
+  /// Examples where this helps:
+  /// * approximate multiples of 0.01
+  /// * approximate multiples of pi
+  ///
+  /// Float mults can work even when there are NaNs and infinities.
+  /// When this is helpful, compression and decompression speeds can be
+  /// substantially reduced. In rare cases, this configuration
+  /// may reduce compression speed somewhat even when it isn't helpful.
+  /// However, the compression ratio improvements tend to be large.
+  ///
+  /// :param float_quant_spec: a FloatQuantSpec that configures whether
+  /// quantized-float detection is enabled.
+  ///
+  /// Examples where this helps:
+  /// * float-valued data stored in a type that is unnecessarily wide (e.g.
+  /// stored as `f64`s where only a `f32` worth of precision is used)
+  ///
   /// :param paging_spec: a PagingSpec describing how many numbers should
   /// go into each page.
   ///
@@ -106,15 +215,17 @@ impl PyChunkConfig {
   #[pyo3(signature = (
     compression_level=pco::DEFAULT_COMPRESSION_LEVEL,
     delta_encoding_order=None,
-    int_mult_spec="enabled".to_string(),
-    float_mult_spec="enabled".to_string(),
+    int_mult_spec=PyIntMultSpec::default(),
+    float_mult_spec=PyFloatMultSpec::default(),
+    float_quant_spec=PyFloatQuantSpec::default(),
     paging_spec=PyPagingSpec::default(),
   ))]
   fn new(
     compression_level: usize,
     delta_encoding_order: Option<usize>,
-    int_mult_spec: String,
-    float_mult_spec: String,
+    int_mult_spec: PyIntMultSpec,
+    float_mult_spec: PyFloatMultSpec,
+    float_quant_spec: PyFloatQuantSpec,
     paging_spec: PyPagingSpec,
   ) -> Self {
     Self {
@@ -122,6 +233,7 @@ impl PyChunkConfig {
       delta_encoding_order,
       int_mult_spec,
       float_mult_spec,
+      float_quant_spec,
       paging_spec,
     }
   }
@@ -131,31 +243,12 @@ impl TryFrom<&PyChunkConfig> for ChunkConfig {
   type Error = PyErr;
 
   fn try_from(py_config: &PyChunkConfig) -> Result<Self, Self::Error> {
-    let int_mult_spec = match py_config.int_mult_spec.to_lowercase().as_str() {
-      "enabled" => IntMultSpec::Enabled,
-      "disabled" => IntMultSpec::Disabled,
-      other => {
-        return Err(PyRuntimeError::new_err(format!(
-          "unknown int mult spec: {}",
-          other
-        )))
-      }
-    };
-    let float_mult_spec = match py_config.float_mult_spec.to_lowercase().as_str() {
-      "enabled" => FloatMultSpec::Enabled,
-      "disabled" => FloatMultSpec::Disabled,
-      other => {
-        return Err(PyRuntimeError::new_err(format!(
-          "unknown float mult spec: {}",
-          other
-        )))
-      }
-    };
     let res = ChunkConfig::default()
       .with_compression_level(py_config.compression_level)
       .with_delta_encoding_order(py_config.delta_encoding_order)
-      .with_int_mult_spec(int_mult_spec)
-      .with_float_mult_spec(float_mult_spec)
+      .with_int_mult_spec(py_config.int_mult_spec.0)
+      .with_float_mult_spec(py_config.float_mult_spec.0)
+      .with_float_quant_spec(py_config.float_quant_spec.0)
       .with_paging_spec(py_config.paging_spec.0.clone());
     Ok(res)
   }
@@ -182,6 +275,9 @@ pub enum DynTypedPyArrayDyn<'py> {
 fn pcodec(py: Python, m: &PyModule) -> PyResult<()> {
   m.add("__version__", env!("CARGO_PKG_VERSION"))?;
   m.add_class::<PyProgress>()?;
+  m.add_class::<PyIntMultSpec>()?;
+  m.add_class::<PyFloatMultSpec>()?;
+  m.add_class::<PyFloatQuantSpec>()?;
   m.add_class::<PyPagingSpec>()?;
   m.add_class::<PyChunkConfig>()?;
   m.add(


Test patch:
diff --git a/pco_python/test/test_standalone.py b/pco_python/test/test_standalone.py
index fde9496b..a60e7d21 100644
--- a/pco_python/test/test_standalone.py
+++ b/pco_python/test/test_standalone.py
@@ -1,5 +1,5 @@
 import numpy as np
-from pcodec import standalone, ChunkConfig, PagingSpec
+from pcodec import ChunkConfig, FloatMultSpec, FloatQuantSpec, IntMultSpec, PagingSpec, standalone
 import pytest
 
 np.random.seed(12345)
@@ -10,7 +10,28 @@
   [9, 100],
 )
 
-all_dtypes = ('f2', 'f4', 'f8', 'i2', 'i4', 'i8', 'u2', 'u4', 'u8')
+all_dtypes = ("f2", "f4", "f8", "i2", "i4", "i8", "u2", "u4", "u8")
+
+all_int_mult_specs = (
+    IntMultSpec.enabled(),
+    IntMultSpec.disabled(),
+    IntMultSpec.provided(2),
+    IntMultSpec.provided(10),
+)
+
+all_float_mult_specs = (
+    FloatMultSpec.enabled(),
+    FloatMultSpec.disabled(),
+    FloatMultSpec.provided(np.pi),
+    FloatMultSpec.provided(10 * np.pi),
+)
+
+all_float_quant_specs = (
+    FloatQuantSpec.disabled(),
+    FloatQuantSpec.provided(4),
+    FloatQuantSpec.provided(16),
+)
+
 
 @pytest.mark.parametrize("shape", all_shapes)
 @pytest.mark.parametrize("dtype", all_dtypes)
@@ -92,13 +113,66 @@ def test_compression_options():
 
   # this is mostly just to check that there is no error, but these settings
   # should give worse compression than the defaults
-  assert len(standalone.simple_compress(
+  assert (len(
+      standalone.simple_compress(
+          data,
+          ChunkConfig(
+              compression_level=0,
+              delta_encoding_order=1,
+              int_mult_spec=IntMultSpec.disabled(),
+              float_mult_spec=FloatMultSpec.disabled(),
+              float_quant_spec=FloatQuantSpec.disabled(),
+              paging_spec=PagingSpec.equal_pages_up_to(77),
+          ),
+      )) > default_size)
+
+
+@pytest.mark.parametrize("int_mult_spec", all_int_mult_specs)
+def test_compression_int_mult_spec_options(int_mult_spec):
+  data = (np.random.normal(size=100) * 1000).astype(np.int32)
+
+  # check for errors
+  compressed = standalone.simple_compress(
+    data,
+    ChunkConfig(int_mult_spec=int_mult_spec),
+  )
+
+  out = standalone.simple_decompress(compressed)
+
+  # check that the decompressed data is correct
+  np.testing.assert_array_equal(data, out)
+
+
+@pytest.mark.parametrize("float_mult_spec", all_float_mult_specs)
+def test_compression_float_mult_spec_options(float_mult_spec):
+  data = (np.random.normal(size=100) * 1000).astype(np.int32) * np.pi
+
+  # check for errors
+  compressed = standalone.simple_compress(
+    data,
+    ChunkConfig(float_mult_spec=float_mult_spec),
+  )
+
+  out = standalone.simple_decompress(compressed)
+
+  # check that the decompressed data is correct
+  np.testing.assert_array_equal(data, out)
+
+
+@pytest.mark.parametrize("float_quant_spec", all_float_quant_specs)
+def test_compression_float_quant_spec_options(float_quant_spec):
+  data = np.random.normal(size=100).astype(np.float32).astype(np.float64)
+
+  # check for errors
+  compressed = standalone.simple_compress(
     data,
     ChunkConfig(
-      compression_level=0,
-      delta_encoding_order=1,
-      int_mult_spec='disabled',
-      float_mult_spec='DISABLED',
-      paging_spec=PagingSpec.equal_pages_up_to(77),
-    )
-  )) > default_size
+      float_mult_spec=FloatMultSpec.disabled(),
+      float_quant_spec=float_quant_spec,
+    ),
+  )
+
+  out = standalone.simple_decompress(compressed)
+
+  # check that the decompressed data is correct
+  np.testing.assert_array_equal(data, out)