cocoindex 0.1.59__tar.gz → 0.1.61__tar.gz

{cocoindex-0.1.59 → cocoindex-0.1.61}/.gitignore

@@ -18,3 +18,5 @@ dist/
 
 # Output of `cocoindex eval`
 examples/**/eval_*
+
+/.vscode

{cocoindex-0.1.59 → cocoindex-0.1.61}/Cargo.lock

@@ -1040,7 +1040,7 @@ dependencies = [
 
 [[package]]
 name = "cocoindex"
-version = "0.1.59"
+version = "0.1.61"
 dependencies = [
  "anyhow",
  "async-openai",

{cocoindex-0.1.59 → cocoindex-0.1.61}/Cargo.toml

@@ -2,7 +2,7 @@
 name = "cocoindex"
 # Version used for local development is always higher than others to take precedence.
 # Will be overridden for specific release versions.
-version = "0.1.59"
+version = "0.1.61"
 edition = "2024"
 rust-version = "1.86"
 

{cocoindex-0.1.59 → cocoindex-0.1.61}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: cocoindex
-Version: 0.1.59
+Version: 0.1.61
 Requires-Dist: click>=8.1.8
 Requires-Dist: rich>=14.0.0
 Requires-Dist: python-dotenv>=1.1.0

{cocoindex-0.1.59 → cocoindex-0.1.61}/docs/docs/core/cli.mdx

@@ -47,12 +47,6 @@ CocoIndex CLI supports the following global options:
 * `--version`: Show the CocoIndex version and exit.
 * `--help`: Show the main help message and exit.
 
-:::caution Deprecated Usage
-
-The old method of invoking the CLI using `python main.py cocoindex ...` via the `@cocoindex.main_fn()` decorator is now deprecated. Please remove `@cocoindex.main_fn()` from your scripts and use the standalone cocoindex command as described.
-
-:::
-
 ## Subcommands
 
 The following subcommands are available:

cocoindex-0.1.61/docs/docs/core/data_types.mdx

@@ -0,0 +1,188 @@
+---
+title: Data Types
+description: Data Types in CocoIndex
+toc_max_heading_level: 4
+---
+
+# Data Types in CocoIndex
+
+In CocoIndex, all data processed by the flow have a type determined when the flow is defined, before any actual data is processed at runtime.
+
+This makes schema of data processed by CocoIndex clear, and easily determine the schema of your index.
+
+## Data Types
+
+As an engine written in Rust, designed to be used in different languages and data are always serializable, CocoIndex defines a type system independent of any specific programming language.
+
+CocoIndex automatically infers data types of the output created by CocoIndex sources and functions.
+You don't need to spell out any data type explicitly when you define the flow.
+All you need to do is to make sure the data passed to functions and targets are compatible with them.
+
+Each type in CocoIndex type system is mapped to one or multiple types in Python.
+When you define a [custom function](/docs/core/custom_function), you need to annotate the data types of arguments and return values.
+
+*   For return values, type annotation is required. Because this provides the ground truth to define the type of the output of the custom function.
+*   For arguments, type annotation is only used to enable the conversion from data values already existing in CocoIndex engine to Python value.
+    Type annotation is optional for basic types. When not specified, CocoIndex will use the *default Python type* for the argument.
+    Type annotation is required for arguments of struct types and table types.
+
+### Basic Types
+
+#### Primitive Types
+
+Primitive types are basic types that are not composed of other types.
+This is the list of all primitive types supported by CocoIndex:
+
+| CocoIndex Type | Python Types | Convertible to | Explanation |
+|------|-------------|--------------|----------------|
+| *Bytes* | `bytes` | | |
+| *Str* | `str` | | |
+| *Bool* | `bool` | | |
+| *Int64* | `cocoindex.Int64`, `int`, `numpy.int64` | | |
+| *Float32* | `cocoindex.Float32`, `numpy.float32` | *Float64* | |
+| *Float64* |  `cocoindex.Float64`, `float`, `numpy.float64` | | |
+| *Range* | `cocoindex.Range` | | |
+| *Uuid* | `uuid.UUId` | | |
+| *Date* | `datetime.date` | | |
+| *Time* | `datetime.time` | | |
+| *LocalDatetime* | `cocoindex.LocalDateTime` | *OffsetDatetime* | without timezone |
+| *OffsetDatetime* | `cocoindex.OffsetDateTime`, `datetime.datetime` | | with timezone |
+| *TimeDelta* | `datetime.timedelta` | | |
+
+Notes:
+
+*   For some CocoIndex types, we support multiple Python types. You can annotate with any of these Python types.
+    The first one is the *default Python type*, which means CocoIndex will create a value with this type when you don't annotate the type in function arguments.
+
+*   All Python types starting with `cocoindex.` are type aliases exported by CocoIndex. They're annotated types based on certain Python types:
+
+    *   `cocoindex.Int64`: `int`
+    *   `cocoindex.Float64`: `float`
+    *   `cocoindex.Float32`: `float`
+    *   `cocoindex.Range`: `tuple[int, int]`, i.e. a start offset (inclusive) and an end offset (exclusive)
+    *   `cocoindex.OffsetDateTime`: `datetime.datetime`
+    *   `cocoindex.LocalDateTime`: `datetime.datetime`
+
+    These aliases provide a non-ambiguous way to represent a specific type in CocoIndex, given their base Python types can represent a superset of possible values.
+
+*   When we say a CocoIndex type is *convertible to* another type, it means Python types for the second type can be also used to bind to a value of the first type.
+    For example, *Float32* is convertible to *Float64*, so you can bind a value of *Float32* to a Python value of `float` or `np.float64` types.
+    For *LocalDatetime*, when you use `cocoindex.OffsetDateTime` or `datetime.datetime` as the annotation to bind its value, the timezone will be set to UTC.
+
+
+#### Json Type
+
+*Json* type can hold any data convertible to JSON by `json` package.
+In Python, it's represented by `cocoindex.Json`.
+It's useful to hold data without fixed schema known at flow definition time.
+
+
+#### Vector Types
+
+A vector type is a collection of elements of the same basic type.
+Optionally, it can have a fixed dimension. Noted as *Vector[Type]* or *Vector[Type, Dim]*, e.g. *Vector[Float32]* or *Vector[Float32, 384]*.
+
+It supports the following Python types:
+
+*   `cocoindex.Vector[T]` or `cocoindex.Vector[T, typing.Literal[Dim]]`, e.g. `cocoindex.Vector[cocoindex.Float32]` or `cocoindex.Vector[cocoindex.Float32, 384]`
+    *   The underlying Python type is `numpy.typing.NDArray[T]` where `T` is a numpy numeric type (`numpy.int64`, `numpy.float32` or `numpy.float64`), or `list[T]` otherwise
+*   `numpy.typing.NDArray[T]` where `T` is a numpy numeric type
+*   `list[T]`
+
+
+#### Union Types
+
+A union type is a type that can represent values in one of multiple basic types.
+Noted as *Type1* | *Type2* | ..., e.g. *Int64* | *Float32* | *Float64*.
+
+The Python type is `T1 | T2 | ...`, e.g. `cocoindex.Int64 | cocoindex.Float32 | cocoindex.Float64`, `int | float` (equivalent to `cocoindex.Int64 | cocoindex.Float64`)
+
+
+### Struct Types
+
+A *Struct* has a bunch of fields, each with a name and a type.
+
+In Python, a *Struct* type is represented by either a [dataclass](https://docs.python.org/3/library/dataclasses.html)
+or a [NamedTuple](https://docs.python.org/3/library/typing.html#typing.NamedTuple), with all fields annotated with a specific type.
+Both options define a structured type with named fields, but they differ slightly:
+
+- **Dataclass**: A flexible class-based structure, mutable by default, defined using the `@dataclass` decorator.
+- **NamedTuple**: An immutable tuple-based structure, defined using `typing.NamedTuple`.
+
+For example:
+
+```python
+from dataclasses import dataclass
+from typing import NamedTuple
+import datetime
+
+# Using dataclass
+@dataclass
+class Person:
+    first_name: str
+    last_name: str
+    dob: datetime.date
+
+# Using NamedTuple
+class PersonTuple(NamedTuple):
+    first_name: str
+    last_name: str
+    dob: datetime.date
+```
+
+Both `Person` and `PersonTuple` are valid Struct types in CocoIndex, with identical schemas (three fields: `first_name` (Str), `last_name` (Str), `dob` (Date)).
+Choose `dataclass` for mutable objects or when you need additional methods, and `NamedTuple` for immutable, lightweight structures.
+
+### Table Types
+
+A *Table* type models a collection of rows, each with multiple columns.
+Each column of a table has a specific type.
+
+We have two specific types of *Table* types: *KTable* and *LTable*.
+
+#### KTable
+
+*KTable* is a *Table* type whose first column serves as the key.
+The row order of a *KTable* is not preserved.
+Type of the first column (key column) must be a [key type](#key-types).
+
+In Python, a *KTable* type is represented by `dict[K, V]`. 
+The `V` should be a *Struct* type, either a `dataclass` or `NamedTuple`, representing the value fields of each row.
+For example, you can use `dict[str, Person]` or `dict[str, PersonTuple]` to represent a *KTable*, with 4 columns: key (*Str*), `first_name` (*Str*), `last_name` (*Str*), `dob` (*Date*).
+
+Note that if you want to use a *Struct* as the key, you need to ensure its value in Python is immutable. For `dataclass`, annotate it with `@dataclass(frozen=True)`. For `NamedTuple`, immutability is built-in. For example:
+For example:
+
+```python
+@dataclass(frozen=True)
+class PersonKey:
+    id_kind: str
+    id: str
+
+class PersonKeyTuple(NamedTuple):
+    id_kind: str
+    id: str
+```
+
+Then you can use `dict[PersonKey, Person]` or `dict[PersonKeyTuple, PersonTuple]` to represent a KTable keyed by `PersonKey` or `PersonKeyTuple`.
+
+
+#### LTable
+
+*LTable* is a *Table* type whose row order is preserved. *LTable* has no key column.
+
+In Python, a *LTable* type is represented by `list[R]`, where `R` is a dataclass representing a row.
+For example, you can use `list[Person]` to represent a *LTable* with 3 columns: `first_name` (*Str*), `last_name` (*Str*), `dob` (*Date*).
+
+## Key Types
+
+Currently, the following types are key types
+
+- *Bytes*
+- *Str*
+- *Bool*
+- *Int64*
+- *Range*
+- *Uuid*
+- *Date*
+- *Struct* with all fields being key types (using `@dataclass(frozen=True)` or `NamedTuple`)

{cocoindex-0.1.59 → cocoindex-0.1.61}/python/cocoindex/__init__.py

@@ -12,12 +12,21 @@ from .flow import flow_def
 from .flow import EvaluateAndDumpOptions, GeneratedField
 from .flow import FlowLiveUpdater, FlowLiveUpdaterOptions
 from .flow import update_all_flows_async, setup_all_flows, drop_all_flows
-from .lib import init, start_server, stop, main_fn
+from .lib import init, start_server, stop
 from .llm import LlmSpec, LlmApiType
 from .index import VectorSimilarityMetric, VectorIndexDef, IndexOptions
 from .setting import DatabaseConnectionSpec, Settings, ServerSettings
 from .setting import get_app_namespace
-from .typing import Float32, Float64, LocalDateTime, OffsetDateTime, Range, Vector, Json
+from .typing import (
+    Int64,
+    Float32,
+    Float64,
+    LocalDateTime,
+    OffsetDateTime,
+    Range,
+    Vector,
+    Json,
+)
 
 __all__ = [
     # Submodules
@@ -50,7 +59,6 @@ __all__ = [
     "init",
     "start_server",
     "stop",
-    "main_fn",
     # LLM
     "LlmSpec",
     "LlmApiType",
@@ -64,6 +72,7 @@ __all__ = [
     "ServerSettings",
     "get_app_namespace",
     # Typing
+    "Int64",
     "Float32",
     "Float64",
     "LocalDateTime",

{cocoindex-0.1.59 → cocoindex-0.1.61}/python/cocoindex/flow.py

@@ -416,6 +416,11 @@ class _SourceRefreshOptions:
     refresh_interval: datetime.timedelta | None = None
 
 
+@dataclass
+class _ExecutionOptions:
+    max_inflight_count: int | None = None
+
+
 class FlowBuilder:
     """
     A flow builder is used to build a flow.
@@ -439,6 +444,7 @@ class FlowBuilder:
         *,
         name: str | None = None,
         refresh_interval: datetime.timedelta | None = None,
+        max_inflight_count: int | None = None,
     ) -> DataSlice[T]:
         """
         Import a source to the flow.
@@ -454,9 +460,12 @@ class FlowBuilder:
                 self._state.field_name_builder.build_name(
                     name, prefix=_to_snake_case(_spec_kind(spec)) + "_"
                 ),
-                dump_engine_object(
+                refresh_options=dump_engine_object(
                     _SourceRefreshOptions(refresh_interval=refresh_interval)
                 ),
+                execution_options=dump_engine_object(
+                    _ExecutionOptions(max_inflight_count=max_inflight_count)
+                ),
             ),
             name,
         )

cocoindex-0.1.61/python/cocoindex/lib.py

@@ -0,0 +1,32 @@
+"""
+Library level functions and states.
+"""
+
+import warnings
+from typing import Callable, Any
+
+from . import _engine  # type: ignore
+from . import flow, setting
+from .convert import dump_engine_object
+
+
+def init(settings: setting.Settings | None = None) -> None:
+    """
+    Initialize the cocoindex library.
+
+    If the settings are not provided, they are loaded from the environment variables.
+    """
+    settings = settings or setting.Settings.from_env()
+    _engine.init(dump_engine_object(settings))
+    setting.set_app_namespace(settings.app_namespace)
+
+
+def start_server(settings: setting.ServerSettings) -> None:
+    """Start the cocoindex server."""
+    flow.ensure_all_flows_built()
+    _engine.start_server(settings.__dict__)
+
+
+def stop() -> None:
+    """Stop the cocoindex library."""
+    _engine.stop()

{cocoindex-0.1.59 → cocoindex-0.1.61}/python/cocoindex/tests/test_convert.py

@@ -78,16 +78,15 @@ def build_engine_value_decoder(
     return make_engine_value_decoder([], engine_type, python_type or engine_type_in_py)
 
 
-def validate_full_roundtrip(
+def validate_full_roundtrip_to(
     value: Any,
-    value_type: Any = None,
-    *other_decoded_values: tuple[Any, Any],
+    value_type: Any,
+    *decoded_values: tuple[Any, Any],
 ) -> None:
     """
-    Validate the given value doesn't change after encoding, sending to engine (using output_type), receiving back and decoding (using input_type).
+    Validate the given value becomes specific values after encoding, sending to engine (using output_type), receiving back and decoding (using input_type).
 
-    `other_decoded_values` is a tuple of (value, type) pairs.
-    If provided, also validate the value can be decoded to the other types.
+    `decoded_values` is a tuple of (value, type) pairs.
     """
     from cocoindex import _engine  # type: ignore
 
@@ -102,15 +101,27 @@ def validate_full_roundtrip(
     value_from_engine = _engine.testutil.seder_roundtrip(
         encoded_value, encoded_output_type
     )
-    decoder = make_engine_value_decoder([], encoded_output_type, value_type)
-    decoded_value = decoder(value_from_engine)
-    assert eq(decoded_value, value), f"{decoded_value} != {value}"
 
-    if other_decoded_values is not None:
-        for other_value, other_type in other_decoded_values:
-            decoder = make_engine_value_decoder([], encoded_output_type, other_type)
-            other_decoded_value = decoder(value_from_engine)
-            assert eq(other_decoded_value, other_value)
+    for other_value, other_type in decoded_values:
+        decoder = make_engine_value_decoder([], encoded_output_type, other_type)
+        other_decoded_value = decoder(value_from_engine)
+        assert eq(other_decoded_value, other_value)
+
+
+def validate_full_roundtrip(
+    value: Any,
+    value_type: Any,
+    *other_decoded_values: tuple[Any, Any],
+) -> None:
+    """
+    Validate the given value doesn't change after encoding, sending to engine (using output_type), receiving back and decoding (using input_type).
+
+    `other_decoded_values` is a tuple of (value, type) pairs.
+    If provided, also validate the value can be decoded to the other types.
+    """
+    validate_full_roundtrip_to(
+        value, value_type, (value, value_type), *other_decoded_values
+    )
 
 
 def test_encode_engine_value_basic_types() -> None:
@@ -218,17 +229,33 @@ def test_encode_engine_value_none() -> None:
 
 
 def test_roundtrip_basic_types() -> None:
-    validate_full_roundtrip(42, int, (42, None))
-    validate_full_roundtrip(3.25, float, (3.25, Float64))
     validate_full_roundtrip(
-        3.25, Float64, (3.25, float), (np.float64(3.25), np.float64)
+        42, cocoindex.Int64, (42, int), (np.int64(42), np.int64), (42, None)
     )
+    validate_full_roundtrip(42, int, (42, cocoindex.Int64))
+    validate_full_roundtrip(np.int64(42), np.int64, (42, cocoindex.Int64))
+
     validate_full_roundtrip(
-        3.25, Float32, (3.25, float), (np.float32(3.25), np.float32)
+        3.25, Float64, (3.25, float), (np.float64(3.25), np.float64), (3.25, None)
     )
+    validate_full_roundtrip(3.25, float, (3.25, Float64))
+    validate_full_roundtrip(np.float64(3.25), np.float64, (3.25, Float64))
+
+    validate_full_roundtrip(
+        3.25,
+        Float32,
+        (3.25, float),
+        (np.float32(3.25), np.float32),
+        (np.float64(3.25), np.float64),
+        (3.25, Float64),
+        (3.25, None),
+    )
+    validate_full_roundtrip(np.float32(3.25), np.float32, (3.25, Float32))
+
     validate_full_roundtrip("hello", str, ("hello", None))
     validate_full_roundtrip(True, bool, (True, None))
     validate_full_roundtrip(False, bool, (False, None))
+    validate_full_roundtrip((1, 2), cocoindex.Range, ((1, 2), None))
     validate_full_roundtrip(
         datetime.date(2025, 1, 1), datetime.date, (datetime.date(2025, 1, 1), None)
     )
@@ -238,14 +265,37 @@ def test_roundtrip_basic_types() -> None:
         cocoindex.LocalDateTime,
         (datetime.datetime(2025, 1, 2, 3, 4, 5, 123456), datetime.datetime),
     )
+
+    tz = datetime.timezone(datetime.timedelta(hours=5))
     validate_full_roundtrip(
-        datetime.datetime(2025, 1, 2, 3, 4, 5, 123456, datetime.UTC),
+        datetime.datetime(2025, 1, 2, 3, 4, 5, 123456, tz),
+        cocoindex.OffsetDateTime,
+        (
+            datetime.datetime(2025, 1, 2, 3, 4, 5, 123456, tz),
+            datetime.datetime,
+        ),
+    )
+    validate_full_roundtrip(
+        datetime.datetime(2025, 1, 2, 3, 4, 5, 123456, tz),
+        datetime.datetime,
+        (datetime.datetime(2025, 1, 2, 3, 4, 5, 123456, tz), cocoindex.OffsetDateTime),
+    )
+    validate_full_roundtrip_to(
+        datetime.datetime(2025, 1, 2, 3, 4, 5, 123456),
         cocoindex.OffsetDateTime,
         (
             datetime.datetime(2025, 1, 2, 3, 4, 5, 123456, datetime.UTC),
             datetime.datetime,
         ),
     )
+    validate_full_roundtrip_to(
+        datetime.datetime(2025, 1, 2, 3, 4, 5, 123456),
+        datetime.datetime,
+        (
+            datetime.datetime(2025, 1, 2, 3, 4, 5, 123456, datetime.UTC),
+            cocoindex.OffsetDateTime,
+        ),
+    )
 
     uuid_value = uuid.uuid4()
     validate_full_roundtrip(uuid_value, uuid.UUID, (uuid_value, None))

{cocoindex-0.1.59 → cocoindex-0.1.61}/python/cocoindex/typing.py

@@ -40,6 +40,7 @@ class TypeAttr:
 
 Annotation = TypeKind | TypeAttr | VectorInfo
 
+Int64 = Annotated[int, TypeKind("Int64")]
 Float32 = Annotated[float, TypeKind("Float32")]
 Float64 = Annotated[float, TypeKind("Float64")]
 Range = Annotated[tuple[int, int], TypeKind("Range")]

{cocoindex-0.1.59 → cocoindex-0.1.61}/src/base/spec.rs

@@ -253,6 +253,11 @@ impl SpecFormatter for OpSpec {
     }
 }
 
+#[derive(Debug, Clone, Serialize, Deserialize, Default)]
+pub struct ExecutionOptions {
+    pub max_inflight_count: Option<u32>,
+}
+
 #[derive(Debug, Clone, Serialize, Deserialize, Default)]
 pub struct SourceRefreshOptions {
     pub refresh_interval: Option<std::time::Duration>,
@@ -274,6 +279,9 @@ pub struct ImportOpSpec {
 
     #[serde(default)]
     pub refresh_options: SourceRefreshOptions,
+
+    #[serde(default)]
+    pub execution_options: ExecutionOptions,
 }
 
 impl SpecFormatter for ImportOpSpec {

{cocoindex-0.1.59 → cocoindex-0.1.61}/src/builder/analyzer.rs

@@ -697,6 +697,9 @@ impl AnalyzerContext {
                 primary_key_type,
                 name: op_name,
                 refresh_options: import_op.spec.refresh_options,
+                concurrency_controller: utils::ConcurrencyController::new(
+                    import_op.spec.execution_options.max_inflight_count,
+                ),
             })
         };
         Ok(result_fut)

{cocoindex-0.1.59 → cocoindex-0.1.61}/src/builder/flow_builder.rs

@@ -288,7 +288,7 @@ impl FlowBuilder {
         OpScopeRef(self.root_op_scope.clone())
     }
 
-    #[pyo3(signature = (kind, op_spec, target_scope, name, refresh_options=None))]
+    #[pyo3(signature = (kind, op_spec, target_scope, name, refresh_options=None, execution_options=None))]
     pub fn add_source(
         &mut self,
         py: Python<'_>,
@@ -297,6 +297,7 @@ impl FlowBuilder {
         target_scope: Option<OpScopeRef>,
         name: String,
         refresh_options: Option<py::Pythonized<spec::SourceRefreshOptions>>,
+        execution_options: Option<py::Pythonized<spec::ExecutionOptions>>,
     ) -> PyResult<DataSlice> {
         if let Some(target_scope) = target_scope {
             if *target_scope != self.root_op_scope {
@@ -313,6 +314,9 @@ impl FlowBuilder {
                     spec: op_spec.into_inner(),
                 },
                 refresh_options: refresh_options.map(|o| o.into_inner()).unwrap_or_default(),
+                execution_options: execution_options
+                    .map(|o| o.into_inner())
+                    .unwrap_or_default(),
             },
         };
         let analyzer_ctx = AnalyzerContext {

{cocoindex-0.1.59 → cocoindex-0.1.61}/src/builder/plan.rs

@@ -56,6 +56,7 @@ pub struct AnalyzedImportOp {
     pub output: AnalyzedOpOutput,
     pub primary_key_type: schema::ValueType,
     pub refresh_options: spec::SourceRefreshOptions,
+    pub concurrency_controller: utils::ConcurrencyController,
 }
 
 pub struct AnalyzedFunctionExecInfo {

{cocoindex-0.1.59 → cocoindex-0.1.61}/src/execution/source_indexer.rs

@@ -282,6 +282,7 @@ impl SourceIndexingContext {
             state.scan_generation
         };
         while let Some(row) = rows_stream.next().await {
+            let _ = import_op.concurrency_controller.acquire().await?;
             for row in row? {
                 self.process_source_key_if_newer(
                     row.key,

{cocoindex-0.1.59 → cocoindex-0.1.61}/src/ops/functions/split_recursively.rs

@@ -918,16 +918,14 @@ impl SimpleFunctionExecutor for Executor {
         let table = output
             .into_iter()
             .map(|chunk_output| {
+                let output_start = chunk_output.start_pos.output.unwrap();
+                let output_end = chunk_output.end_pos.output.unwrap();
                 (
-                    RangeValue::new(
-                        chunk_output.start_pos.byte_offset,
-                        chunk_output.end_pos.byte_offset,
-                    )
-                    .into(),
+                    RangeValue::new(output_start.char_offset, output_end.char_offset).into(),
                     fields_value!(
                         Arc::<str>::from(chunk_output.text),
-                        chunk_output.start_pos.output.unwrap().into_output(),
-                        chunk_output.end_pos.output.unwrap().into_output()
+                        output_start.into_output(),
+                        output_end.into_output()
                     )
                     .into(),
                 )

{cocoindex-0.1.59 → cocoindex-0.1.61}/src/py/convert.rs

@@ -1,3 +1,5 @@
+use crate::prelude::*;
+
 use bytes::Bytes;
 use numpy::{PyArray1, PyArrayDyn, PyArrayMethods};
 use pyo3::IntoPyObjectExt;
@@ -6,14 +8,10 @@ use pyo3::types::PyAny;
 use pyo3::types::{PyList, PyTuple};
 use pyo3::{exceptions::PyException, prelude::*};
 use pythonize::{depythonize, pythonize};
-use serde::Serialize;
 use serde::de::DeserializeOwned;
-use std::collections::BTreeMap;
 use std::ops::Deref;
-use std::sync::Arc;
 
 use super::IntoPyResult;
-use crate::base::{schema, value};
 
 #[derive(Debug)]
 pub struct Pythonized<T>(pub T);
@@ -143,7 +141,23 @@ fn basic_value_from_py_object<'py>(
             value::BasicValue::LocalDateTime(v.extract::<chrono::NaiveDateTime>()?)
         }
         schema::BasicValueType::OffsetDateTime => {
-            value::BasicValue::OffsetDateTime(v.extract::<chrono::DateTime<chrono::FixedOffset>>()?)
+            if v.getattr_opt("tzinfo")?
+                .ok_or_else(|| {
+                    PyErr::new::<PyTypeError, _>(format!(
+                        "expecting a datetime.datetime value, got {}",
+                        v.get_type()
+                    ))
+                })?
+                .is_none()
+            {
+                value::BasicValue::OffsetDateTime(
+                    v.extract::<chrono::NaiveDateTime>()?.and_utc().into(),
+                )
+            } else {
+                value::BasicValue::OffsetDateTime(
+                    v.extract::<chrono::DateTime<chrono::FixedOffset>>()?,
+                )
+            }
         }
         schema::BasicValueType::TimeDelta => {
             value::BasicValue::TimeDelta(v.extract::<chrono::TimeDelta>()?)

cocoindex-0.1.61/src/utils/concur_control.rs

@@ -0,0 +1,30 @@
+use crate::prelude::*;
+
+use tokio::sync::{Semaphore, SemaphorePermit};
+
+pub struct ConcurrencyController {
+    inflight_count_sem: Option<Semaphore>,
+}
+
+pub struct ConcurrencyControllerPermit<'a> {
+    _inflight_count_permit: Option<SemaphorePermit<'a>>,
+}
+
+impl ConcurrencyController {
+    pub fn new(max_inflight_count: Option<u32>) -> Self {
+        Self {
+            inflight_count_sem: max_inflight_count.map(|max| Semaphore::new(max as usize)),
+        }
+    }
+
+    pub async fn acquire<'a>(&'a self) -> Result<ConcurrencyControllerPermit<'a>> {
+        let inflight_count_permit = if let Some(sem) = &self.inflight_count_sem {
+            Some(sem.acquire().await?)
+        } else {
+            None
+        };
+        Ok(ConcurrencyControllerPermit {
+            _inflight_count_permit: inflight_count_permit,
+        })
+    }
+}

{cocoindex-0.1.59 → cocoindex-0.1.61}/src/utils/mod.rs

@@ -3,3 +3,6 @@ pub mod fingerprint;
 pub mod immutable;
 pub mod retryable;
 pub mod yaml_ser;
+
+mod concur_control;
+pub use concur_control::ConcurrencyController;