yjdata 0.0.1__py3-none-any.whl

yjdata/_optional.py

@@ -0,0 +1,19 @@
+from importlib import import_module
+from types import ModuleType
+
+
+def import_optional_dependency(
+    name: str, *, install: str | None = None,
+) -> ModuleType:
+    """Import a dependency with an optional installation requirement hint."""
+    root = name.split(".", 1)[0]
+    try:
+        return import_module(name)
+    except ModuleNotFoundError as exc:
+        if exc.name != root:
+            raise
+        message = (
+            f"Missing optional dependency '{root}'. "
+            f'Install it with: pip install "{install or root}"'
+        )
+        raise ImportError(message) from exc

yjdata/dataset/__init__.py

@@ -0,0 +1,8 @@
+from .common import (
+    normalize_polars_df,
+    write_chunked_parquet,
+    write_seq_parquet
+)
+from .quote import get_step_quote_fields
+from .snapshot import get_snapshot_l2_fields
+from .trade import get_step_trade_pair_fields, get_step_trade_side_fields

yjdata/dataset/common.py

@@ -0,0 +1,245 @@
+import dataclasses
+import enum
+import json
+import polars as pl
+import pyarrow as pa
+import pyarrow.parquet as pq
+
+from pathlib import Path
+from typing import Any
+
+from ..dtype import DType
+
+
+__all__ = [
+    'DataField',
+    'DataScaleType',
+
+    'normalize_polars_df',
+    'write_chunked_parquet',
+    'write_seq_parquet'
+]
+
+
+_PYARROW_DTYPE_MAP: dict[DType, pa.DataType] = {
+    DType.INT8: pa.int8(),
+    DType.UINT8: pa.uint8(),
+    DType.INT32: pa.int32(),
+    DType.INT64: pa.int64(),
+    DType.FLOAT64: pa.float64(),
+    DType.STRING: pa.string()
+}
+
+_POLARS_DTYPE_MAP: dict[DType, Any] = {
+    DType.INT8: pl.Int8,
+    DType.UINT8: pl.UInt8,
+    DType.INT32: pl.Int32,
+    DType.INT64: pl.Int64,
+    DType.FLOAT64: pl.Float64,
+    DType.STRING: pl.String,
+}
+
+
+class DataScaleType(enum.IntEnum):
+    """Scaling factors for data fields."""
+    NONE = 0
+    PRICE = 1       # scale by 1e6
+    AMOUNT = 2      # scale by 1e3
+
+
+@dataclasses.dataclass
+class DataField:
+    """Define a data column and its conversion rules.
+
+    Parameters
+    ----------
+    name : str
+        Column name.
+    dtype : DType
+        Target data type.
+    allow_empty : bool, default False
+        Whether missing or null values are allowed.
+    default_value : object, optional
+        Value used for missing columns and null entries.
+    scale_type : DataScaleType, default DataScaleType.NONE
+        Scaling rule applied when scale conversion is enabled.
+    """
+    name: str
+    dtype: DType
+    allow_empty: bool           = dataclasses.field(default=False)
+    default_value: Any          = dataclasses.field(default=None)
+    scale_type: DataScaleType   = dataclasses.field(default=DataScaleType.NONE)
+
+    @property
+    def is_int_type(self):
+        return self.dtype in (
+            DType.INT8,
+            DType.UINT8,
+            DType.INT32,
+            DType.INT64,
+        )
+
+    @property
+    def polars_dtype(self):
+        try:
+            return _POLARS_DTYPE_MAP[self.dtype]
+        except KeyError as error:
+            raise ValueError(f"Unsupported dtype '{self.dtype}' for polars conversion.") from error
+
+    @property
+    def pyarrow_dtype(self):
+        try:
+            return _PYARROW_DTYPE_MAP[self.dtype]
+        except KeyError as error:
+            raise ValueError(f"Unsupported dtype '{self.dtype}' for pyarrow conversion.") from error
+
+
+def normalize_polars_df(
+    data: pl.DataFrame,
+    fields: list[DataField],
+    convert_scale: bool = False
+) -> pl.DataFrame:
+    """Normalize column order, types, and optional scaling.
+
+    Parameters
+    ----------
+    data : polars.DataFrame
+        Input rows.
+    fields : list of DataField
+        Target columns in output order. Unlisted columns are omitted.
+    convert_scale : bool, default False
+        Multiply prices by 1e6 and amounts by 1e3 according to scale_type.
+
+    Returns
+    -------
+    polars.DataFrame
+        Selected columns with defaults, scaling, and type conversions applied.
+        Floating-point values are rounded before conversion to integer types.
+
+    Raises
+    ------
+    ValueError
+        If a required column is missing, contains nulls without a default,
+        or has an unsupported target type.
+    """
+    expressions = []
+    column_names = set(data.columns)
+
+    # Track null checks to report the offending column.
+    null_filter_map = {}
+
+    for field in fields:
+        name = field.name
+        polars_dtype = field.polars_dtype
+        expr = pl.col(name)
+
+        if name not in column_names:
+            if not field.allow_empty and field.default_value is None:
+                raise ValueError(f"Missing required column '{name}' which does not allow empty.")
+            expr = pl.lit(field.default_value, dtype=polars_dtype).alias(name)
+
+        column_dtype = data.schema.get(name, polars_dtype)
+        if not field.allow_empty and field.default_value is None:
+            null_filter_map[name] = pl.col(name).is_null()
+
+        if field.default_value is not None:
+            expr = expr.fill_null(field.default_value)
+
+        if convert_scale:
+            if field.scale_type == DataScaleType.PRICE:
+                expr = expr * 1_000_000
+            elif field.scale_type == DataScaleType.AMOUNT:
+                expr = expr * 1_000
+
+        if column_dtype in (pl.Float32, pl.Float64) and field.is_int_type:
+            expr = expr.round(0).cast(polars_dtype, strict=False)
+
+        expressions.append(expr.cast(polars_dtype).alias(name))
+
+    # Check required columns for null values.
+    null_filter = pl.lit(False)
+    for f in null_filter_map.values():
+        null_filter |= f
+
+    if data.filter(null_filter).height:
+        # Identify the column containing null values.
+        for name, f in null_filter_map.items():
+            if data.filter(f).height:
+                raise ValueError(f'Column "{name}" contains null value, but it is marked as not allowed.')
+
+    return data.select(expressions)
+
+
+def _normalize_string_type(tb: pa.Table) -> pa.Table:
+    columns = []
+    for col in tb.schema:
+        if pa.types.is_large_string(col.type):
+            t = pa.string()
+        else:
+            t = col.type
+        columns.append(pa.field(col.name, t))
+    schema = pa.schema(columns)
+    return tb.cast(schema)
+
+
+def write_chunked_parquet(
+        df: pl.DataFrame,
+        path: str | Path,
+        chunked_by: str = 'Symbol',
+        sort_chunks: bool = True,
+        metadata_key: bytes = b'group_symbols'
+):
+    """Write grouped rows to a single Parquet file.
+
+    Parameters
+    ----------
+    df : polars.DataFrame
+        Rows to write.
+    path : str or pathlib.Path
+        Output file path.
+    chunked_by : str, default 'Symbol'
+        Column used to group rows.
+    sort_chunks : bool, default True
+        Sort groups by key. Otherwise, use the order returned by Polars.
+    metadata_key : bytes, default b'group_symbols'
+        Schema metadata key for a JSON mapping from stringified group keys
+        to zero-based group indices.
+
+    Notes
+    -----
+    Each group is written in a separate write call. Arrow large string
+    columns are converted to string columns before writing.
+    """
+    groups = list(df.group_by(chunked_by))
+    if sort_chunks:
+        groups.sort(key=lambda g: g[0])
+
+    base_table = _normalize_string_type(df.head(0).to_arrow())
+    metadata = dict(base_table.schema.metadata or {})
+    metadata[metadata_key] = ','.join(str(key) for key, _ in groups).encode('utf-8')
+
+    with pq.ParquetWriter(path, base_table.schema.with_metadata(metadata)) as writer:
+        for _, group in groups:
+            table = _normalize_string_type(group.to_arrow())
+            writer.write_table(table)
+
+
+def write_seq_parquet(
+        df: pl.DataFrame,
+        path: str | Path,
+):
+    """Write rows to a Parquet file in their existing order.
+
+    Parameters
+    ----------
+    df : polars.DataFrame
+        Rows to write without additional grouping or sorting.
+    path : str or pathlib.Path
+        Output file path.
+
+    Notes
+    -----
+    Arrow large string columns are converted to string columns before writing.
+    """
+    table = _normalize_string_type(df.to_arrow())
+    pq.write_table(table, path)

yjdata/dataset/quote.py

@@ -0,0 +1,30 @@
+from typing import List
+
+from .common import DataField, DataScaleType
+from ..dtype import DType
+
+
+__all__ = [
+    'get_step_quote_fields',
+]
+
+
+def get_step_quote_fields() -> List[DataField]:
+    return [
+        DataField('Symbol',         DType.STRING),
+        DataField('Exchange',       DType.STRING),
+        DataField('TradeDate',      DType.INT32),
+        DataField('Time',           DType.INT64),
+        DataField('RecvTime',       DType.INT64, True),
+        DataField('Price',          DType.INT64, scale_type=DataScaleType.PRICE),
+        DataField('Qty',            DType.INT64),
+        DataField('ModType',        DType.UINT8),
+        DataField('PriceType',      DType.UINT8, False, 0),
+        DataField('OrderType',      DType.UINT8, False, 0),
+        DataField('Side',           DType.UINT8),
+        DataField('PositionOffset', DType.UINT8, True),
+        DataField('ChannelId',      DType.INT64, False, 0),
+        DataField('OrderNo',        DType.INT64, False, 0),
+        DataField('RecNo',          DType.INT64, False, 0),
+        DataField('MessageNo',      DType.INT64, False, 0),
+    ]

yjdata/dataset/snapshot.py

@@ -0,0 +1,51 @@
+from typing import List
+
+from .common import DataField, DataScaleType
+from ..dtype import DType
+
+
+__all__ = [
+    'get_snapshot_l2_fields'
+]
+
+
+def get_snapshot_l2_fields() -> List[DataField]:
+    required_level_fields = []
+    for prefix in ['Bid', 'BidQty', 'Ask', 'AskQty']:
+        for i in range(1, 11):
+            if prefix.endswith('Qty'):
+                required_level_fields.append(DataField(f'{prefix}{i}', DType.INT64, True, 0))
+            else:
+                required_level_fields.append(DataField(f'{prefix}{i}', DType.INT64, True, scale_type=DataScaleType.PRICE))
+
+    return [
+        DataField('Symbol',     DType.STRING),
+        DataField('Exchange',   DType.STRING),
+        DataField('TradeDate',  DType.INT32),
+        DataField('Time',       DType.INT64),
+        DataField('RecvTime',   DType.INT64, True),
+        DataField('Phase',      DType.UINT8, True),
+        DataField('PreClose',   DType.INT64, True, scale_type=DataScaleType.PRICE),
+        DataField('PreSettle',  DType.INT64, True, scale_type=DataScaleType.PRICE),
+        DataField('PreOpenInterest', DType.INT64, True),
+        DataField('Open',       DType.INT64, True, scale_type=DataScaleType.PRICE),
+        DataField('High',       DType.INT64, True, scale_type=DataScaleType.PRICE),
+        DataField('Low',        DType.INT64, True, scale_type=DataScaleType.PRICE),
+        DataField('Last',       DType.INT64, scale_type=DataScaleType.PRICE),
+        DataField('Qty',        DType.INT64),
+        DataField('Amount',     DType.INT64, scale_type=DataScaleType.AMOUNT),
+        *required_level_fields,
+        DataField('UpperLimit', DType.INT64, True, scale_type=DataScaleType.PRICE),
+        DataField('LowerLimit', DType.INT64, True, scale_type=DataScaleType.PRICE),
+        DataField('OpenInterest', DType.INT64, True),
+        DataField('TradeCount', DType.INT64, True),
+        DataField('BidQty',     DType.INT64, True),
+        DataField('AskQty',     DType.INT64, True),
+        DataField('BidCount',   DType.INT64, True),
+        DataField('AskCount',   DType.INT64, True),
+        DataField('BidAvg',     DType.INT64, True, scale_type=DataScaleType.PRICE),
+        DataField('AskAvg',     DType.INT64, True, scale_type=DataScaleType.PRICE),
+        DataField('Close',      DType.INT64, True, scale_type=DataScaleType.PRICE),
+        DataField('Settle',     DType.INT64, True, scale_type=DataScaleType.PRICE),
+        DataField('MessageNo',  DType.INT64, False, 0),
+    ]

yjdata/dataset/trade.py

@@ -0,0 +1,45 @@
+from typing import List
+
+from .common import DataField, DataScaleType
+from ..dtype import DType
+
+
+__all__ = [
+    'get_step_trade_side_fields',
+    'get_step_trade_pair_fields'
+]
+
+
+def get_step_trade_side_fields() -> List[DataField]:
+    return [
+        DataField('Symbol',     DType.STRING),
+        DataField('Exchange',   DType.STRING),
+        DataField('TradeDate',  DType.INT32),
+        DataField('Time',       DType.INT64),
+        DataField('RecvTime',   DType.INT64, True),
+        DataField('Price',      DType.INT64, scale_type=DataScaleType.PRICE),
+        DataField('Qty',        DType.INT64),
+        DataField('OrderId',    DType.INT64, False, 0),
+        DataField('Side',       DType.UINT8),
+        DataField('PairNo',     DType.INT64, False, 0),
+        DataField('ChannelId',  DType.INT64, False, 0),
+        DataField('RecNo',      DType.INT64, False, 0),
+        DataField('MessageNo',  DType.INT64, False, 0),
+    ]
+
+
+def get_step_trade_pair_fields() -> List[DataField]:
+    return [
+        DataField('Symbol',     DType.STRING),
+        DataField('Exchange',   DType.STRING),
+        DataField('TradeDate',  DType.INT32),
+        DataField('Time',       DType.INT64),
+        DataField('RecvTime',   DType.INT64, True),
+        DataField('Price',      DType.INT64, scale_type=DataScaleType.PRICE),
+        DataField('Qty',        DType.INT64),
+        DataField('BidOrderId', DType.INT64, False, 0),
+        DataField('AskOrderId', DType.INT64, False, 0),
+        DataField('ChannelId',  DType.INT64, False, 0),
+        DataField('RecNo',      DType.INT64, False, 0),
+        DataField('MessageNo',  DType.INT64, False, 0),
+    ]

yjdata/db/__init__.py

@@ -0,0 +1,15 @@
+from .adapter import DbAdapter
+from .expression import Expression, col, lit
+from .sql import SqlAdapter
+from .mysql import MysqlAdapter
+from .postgres import PostgresAdapter
+
+__all__ = [
+    "DbAdapter",
+    "SqlAdapter",
+    "Expression",
+    "col",
+    "lit",
+    "PostgresAdapter",
+    "MysqlAdapter"
+]

yjdata/db/adapter.py

@@ -0,0 +1,250 @@
+import polars as pl
+
+from abc import ABC, abstractmethod
+from typing import Any, Sequence
+
+from .expression import Expression
+
+
+def validate_batch_size(batch_size: int) -> None:
+    """Require a positive integer batch size, excluding booleans."""
+    if isinstance(batch_size, bool) or not isinstance(batch_size, int) or batch_size <= 0:
+        raise ValueError("batch_size must be a positive integer")
+
+
+def validate_insert_args(df: pl.DataFrame, batch_size: int) -> None:
+    """Validate batch size and require at least one DataFrame column."""
+    validate_batch_size(batch_size)
+    if not df.columns:
+        raise ValueError("Cannot insert a DataFrame with no columns")
+
+
+def validate_delete_args(filter: Expression | None) -> None:
+    """Accept a deletion expression or None."""
+    if filter is not None and not isinstance(filter, Expression):
+        raise TypeError("filter must be an Expression or None")
+
+
+def validate_select_args(filter: Expression | None, columns: Sequence[str] | None) -> None:
+    """Require an optional expression and a nonempty sequence of column names."""
+    if filter is not None and not isinstance(filter, Expression):
+        raise TypeError("filter must be an Expression or None")
+    if columns is not None:
+        if isinstance(columns, (str, bytes)) or not isinstance(columns, Sequence):
+            raise TypeError("columns must be a sequence of column names")
+        if not columns:
+            raise ValueError("columns must contain at least one column")
+        if any(not isinstance(column, str) for column in columns):
+            raise TypeError("Column names must be strings")
+        if any(not column or "\x00" in column for column in columns):
+            raise ValueError("Column names must be nonempty and contain no NUL characters")
+
+
+def validate_update_args(
+    keys: Sequence[str],
+    df: pl.DataFrame,
+    batch_size: int,
+    *,
+    require_update_columns: bool = True,
+) -> list[str]:
+    """Validate update/upsert arguments and return non-key columns in input order."""
+    validate_batch_size(batch_size)
+    if not keys:
+        raise ValueError("keys must contain at least one column")
+
+    missing_keys = [key for key in keys if key not in df.columns]
+    if missing_keys:
+        raise ValueError(f"Key columns are missing from the DataFrame: {missing_keys}")
+
+    update_columns = [column for column in df.columns if column not in keys]
+    if require_update_columns and not update_columns:
+        raise ValueError("The DataFrame must contain at least one non-key column")
+    return update_columns
+
+
+class DbAdapter(ABC):
+    """Interface for exchanging tabular data between databases and Python objects."""
+
+    @abstractmethod
+    def delete_data(self, table: str, filter: Expression | None = None) -> None:
+        """Delete rows matching an expression.
+
+        Parameters
+        ----------
+        table : str
+            Target table name, optionally qualified by schema or database.
+        filter : Expression or None, default None
+            Condition identifying rows to delete. None deletes all rows.
+
+        Raises
+        ------
+        TypeError
+            If filter is neither an Expression nor None.
+        """
+
+    def insert_data(self, table: str, data: list[dict[str, Any]], batch_size: int = 65536) -> None:
+        """Insert dictionary rows in batches within one transaction.
+
+        Parameters
+        ----------
+        table : str
+            Target table name, optionally qualified by schema or database.
+        data : list of dict
+            Rows to insert. An empty list returns without validation.
+        batch_size : int, default 65536
+            Maximum rows per batch. Must be a positive integer.
+
+        Raises
+        ------
+        ValueError
+            If nonempty data has no columns or batch_size is invalid.
+        """
+        if not data:
+            return
+        df = pl.DataFrame(data)
+        self.insert_polars(table, df, batch_size)
+
+    @abstractmethod
+    def insert_polars(
+        self,
+        table: str,
+        df: pl.DataFrame,
+        batch_size: int = 65536,
+    ) -> None:
+        """Insert DataFrame rows in batches within one transaction.
+
+        Parameters
+        ----------
+        table : str
+            Target table name, optionally qualified by schema or database.
+        df : polars.DataFrame
+            Rows to insert. A frame with columns but no rows is a no-op.
+        batch_size : int, default 65536
+            Maximum rows per batch. Must be a positive integer.
+
+        Raises
+        ------
+        ValueError
+            If batch_size is invalid or df has no columns.
+        """
+
+    def select_data(
+        self,
+        table: str,
+        *,
+        filter: Expression | None = None,
+        columns: Sequence[str] | None = None,
+    ) -> list[dict[str, Any]]:
+        """Select rows from one table as dictionaries.
+
+        Parameters
+        ----------
+        table : str
+            Target table name, optionally qualified by schema or database.
+        filter : Expression, optional
+            Condition to match; None selects all rows.
+        columns : sequence of str, optional
+            Nonempty sequence of column names; None selects all columns.
+
+        Returns
+        -------
+        list of dict
+            Matching rows keyed by column name, without a guaranteed order.
+        """
+        raise NotImplementedError("This adapter does not implement selection")
+
+    def select_polars(
+        self,
+        table: str,
+        *,
+        filter: Expression | None = None,
+        columns: Sequence[str] | None = None,
+    ) -> pl.DataFrame:
+        """Select rows from one table as a Polars DataFrame.
+
+        Parameters
+        ----------
+        table : str
+            Target table name, optionally qualified by schema or database.
+        filter : Expression, optional
+            Condition to match; None selects all rows.
+        columns : sequence of str, optional
+            Nonempty sequence of column names; None selects all columns.
+
+        Returns
+        -------
+        polars.DataFrame
+            Matching rows with types inferred from all rows, without a
+            guaranteed order. An empty result has no columns.
+        """
+        raise NotImplementedError("This adapter does not implement selection")
+
+    @abstractmethod
+    def update_polars(
+        self,
+        table: str,
+        keys: Sequence[str],
+        df: pl.DataFrame,
+        batch_size: int = 65536,
+    ) -> None:
+        """Update non-key columns in rows matching all key columns.
+
+        Parameters
+        ----------
+        table : str
+            Target table name, optionally qualified by schema or database.
+        keys : sequence of str
+            Nonempty list of key columns present in df.
+        df : polars.DataFrame
+            Rows to write. An empty frame is a no-op after validation.
+        batch_size : int, default 65536
+            Maximum rows per batch. Must be a positive integer.
+
+        Raises
+        ------
+        ValueError
+            If batch_size is invalid, keys is empty, a key is missing from df,
+            or df has no non-key columns.
+
+        Notes
+        -----
+        All batches execute within one transaction.
+        """
+
+    @abstractmethod
+    def upsert_polars(
+        self,
+        table: str,
+        keys: Sequence[str],
+        df: pl.DataFrame,
+        batch_size: int = 65536,
+    ) -> None:
+        """Insert rows, updating non-key columns on conflicts.
+
+        Parameters
+        ----------
+        table : str
+            Target table name, optionally qualified by schema or database.
+        keys : sequence of str
+            Nonempty list of key columns present in df.
+        df : polars.DataFrame
+            Rows to write. An empty frame is a no-op after validation.
+        batch_size : int, default 65536
+            Maximum rows per batch. Must be a positive integer.
+
+        Raises
+        ------
+        ValueError
+            If batch_size is invalid, keys is empty, or a key is missing from df.
+
+        Notes
+        -----
+        All batches execute within one transaction.
+
+        PostgreSQL requires keys to match a primary key or unique constraint
+        usable with ON CONFLICT. MySQL updates on any primary key or unique
+        constraint conflict, regardless of keys.
+
+        If df contains only key columns, PostgreSQL skips conflicts with
+        DO NOTHING; MySQL assigns a key column to itself to preserve values.
+        """