yjdata 0.0.1__tar.gz

yjdata-0.0.1/PKG-INFO

@@ -0,0 +1,8 @@
+Metadata-Version: 2.4
+Name: yjdata
+Version: 0.0.1
+Summary: yjdata SDK for Python
+Requires-Python: >=3.10
+Description-Content-Type: text/markdown
+Requires-Dist: polars>=1.30
+Requires-Dist: pyarrow>=25.0.1

yjdata-0.0.1/pyproject.toml

@@ -0,0 +1,17 @@
+[build-system]
+requires = ["setuptools>=70.0.0", "wheel>=0.36.0"]
+build-backend = "setuptools.build_meta"
+
+[project]
+name = "yjdata"
+version = "0.0.1"
+description = "yjdata SDK for Python"
+readme = "README.md"
+requires-python = ">=3.10"
+dependencies = [
+    "polars>=1.30",
+    "pyarrow>=25.0.1",
+]
+
+[tool.setuptools.packages.find]
+where = ["src/"]

yjdata-0.0.1/setup.cfg

@@ -0,0 +1,4 @@
+[egg_info]
+tag_build = 
+tag_date = 0
+

yjdata-0.0.1/src/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-0.0.1/src/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-0.0.1/src/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-0.0.1/src/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-0.0.1/src/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-0.0.1/src/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-0.0.1/src/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"
+]